diff --git a/src/cmap/_colormap.py b/src/cmap/_colormap.py index b681c9a11..4dab625fe 100644 --- a/src/cmap/_colormap.py +++ b/src/cmap/_colormap.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: from collections.abc import Iterator - from typing import Literal, TypeAlias, TypeGuard + from typing import Literal, NotRequired, TypeAlias, TypeGuard import bokeh.models import matplotlib.colors @@ -50,7 +50,21 @@ class ColormapDict(TypedDict): identifier: str category: str | None value: list[tuple[float, list[float]]] - + # written only when they hold non-default state, so a colormap that has none + # of them serializes exactly as it did before they existed + interpolation: NotRequired[Interpolation] + under: NotRequired[list[float]] + over: NotRequired[list[float]] + bad: NotRequired[list[float]] + neg_inf: NotRequired[list[float]] + pos_inf: NotRequired[list[float]] + nan: NotRequired[list[float]] + + +# the extreme colors, in constructor-argument spelling; `_color` is the property. +# `as_dict()` and pickle iterate this; `reversed()`, `shifted()` and `with_extremes()` +# spell out each one, so a new extreme must be added there too. +_EXTREME_FIELDS = ("under", "over", "bad", "neg_inf", "pos_inf", "nan") LutCallable: TypeAlias = Callable[["NDArray"], "NDArray"] """Function type for a callable that takes an array of values in the range [0, 1] and returns an (N, 4) array of RGBA values in the range [0, 1].""" # noqa @@ -162,6 +176,7 @@ class Colormap: __slots__ = ( "__weakref__", + "_catalog_name", "_has_exceptional", "_initialized", "_lut_cache", @@ -276,14 +291,29 @@ def __init__( cmap_kwargs: dict[str, Any] | None = None, ) -> None: self.info: CatalogItem | None = None + # the catalog name that alone rebuilds this colormap, if nothing overrode it + self._catalog_name: str | None = None if isinstance(value, str): rev = value.endswith("_r") info = self.catalog()[value[:-2] if rev else value] + # every other constructor argument must be listed here: one that is + # missed is silently dropped when this colormap serializes or pickles + # as its bare name + overrides = (name, identifier, category, interpolation, cmap_kwargs) + extremes = (under, over, bad, neg_inf, pos_inf, nan) + if all(v is None for v in (*overrides, *extremes)): + self._catalog_name = info.qualified_name + ("_r" if rev else "") name = name or f"{info.qualified_name}{'_r' if rev else ''}" category = category or info.category - over = info.over if over is None else over - under = info.under if under is None else under + # the record's directional colors follow the ends they extend, so a "_r" + # name swaps them. Swapping here rather than after the lines below keeps + # an explicit `under=`/`over=` argument on the end the caller named. + info_under, info_over = ( + (info.over, info.under) if rev else (info.under, info.over) + ) + over = info_over if over is None else over + under = info_under if under is None else under bad = info.bad if bad is None else bad self.info = info @@ -531,32 +561,64 @@ def with_extremes( pos_inf: ColorLike | None = None, nan: ColorLike | None = None, ) -> Colormap: - """Return a copy of the colormap with new extreme values.""" + """Return a copy of the colormap with new extreme values. + + Colors that are not passed are carried over from this colormap, as in + matplotlib. + """ return type(self)( self.color_stops, name=self.name, + identifier=self.identifier, category=self.category, interpolation=self.interpolation, - bad=bad, - under=under, - over=over, - neg_inf=neg_inf, - pos_inf=pos_inf, - nan=nan, + under=self.under_color if under is None else under, + over=self.over_color if over is None else over, + bad=self.bad_color if bad is None else bad, + neg_inf=self.neg_inf_color if neg_inf is None else neg_inf, + pos_inf=self.pos_inf_color if pos_inf is None else pos_inf, + nan=self.nan_color if nan is None else nan, ) + @property + def _extremes(self) -> dict[str, Color | None]: + """The extreme colors, keyed by their constructor argument name.""" + return {f: getattr(self, f"{f}_color") for f in _EXTREME_FIELDS} + + def _constructor_kwargs(self) -> dict[str, Any]: + """Every keyword argument needed to rebuild this colormap from its stops.""" + kwargs: dict[str, Any] = { + "name": self.name, + "identifier": self.identifier, + "category": self.category, + "interpolation": self.interpolation, + } + kwargs.update({k: v for k, v in self._extremes.items() if v is not None}) + return kwargs + def as_dict(self) -> ColormapDict: """Return a dictionary representation of the colormap. The returned dictionary is suitable for serialization, or for passing to the Colormap constructor. + + Note that a colormap backed by a lut function is sampled into fixed stops + here; use `pickle` or `copy` to duplicate one without resampling. """ - return { + d: ColormapDict = { "name": self.name, "identifier": self.identifier, "category": self.category, "value": [(p, list(c)) for p, c in self.color_stops], } + # optional keys, written only when they hold non-default state: `value` is + # always stops, which the constructor reads as linear. + if self.interpolation != "linear": + d["interpolation"] = self.interpolation + for field, color in self._extremes.items(): + if color is not None: + d[field] = list(color) # type: ignore[literal-required] + return d def lut( self, N: int = 256, gamma: float = 1, *, with_over_under: bool = False @@ -635,6 +697,9 @@ def iter_colors(self, N: Iterable[float] | int | None = None) -> Iterator[Color] def reversed(self, name: str | None = None) -> Colormap: """Return a new Colormap, with reversed colors. + `under`/`over` and `neg_inf`/`pos_inf` are swapped, since they name the ends + they extend. The other extreme colors and the interpolation carry over. + Parameters ---------- name: str | None @@ -646,8 +711,19 @@ def reversed(self, name: str | None = None) -> Colormap: if name is None: name = self.name[:-2] if self.name.endswith("_r") else f"{self.name}_r" + # `identifier` is not carried: it is derived from the name, which changed here return type(self)( - self.color_stops.reversed(), name=name, category=self.category + self.color_stops.reversed(), + name=name, + category=self.category, + interpolation=self.interpolation, + # under/over and neg_inf/pos_inf name the ends they extend, so they follow + under=self.over_color, + over=self.under_color, + neg_inf=self.pos_inf_color, + pos_inf=self.neg_inf_color, + bad=self.bad_color, + nan=self.nan_color, ) def shifted( @@ -740,8 +816,17 @@ def __setattr__(self, _name: str, _value: Any) -> None: object.__setattr__(self, _name, _value) def __reduce__(self) -> str | tuple[Any, ...]: - # for pickle - return self.__class__, (self.color_stops,) + # for pickle. An unmodified catalog colormap goes through as its name, which + # also restores `info`. + if self._catalog_name is not None: + return self.__class__, (self._catalog_name,) + # The stops go through as the object rather than as as_dict()'s samples, so a + # colormap backed by a lut function keeps the function. + return _rebuild_colormap, ( + self.__class__, + self.color_stops, + self._constructor_kwargs(), + ) def __eq__(self, other: object) -> bool: if not isinstance(other, Colormap): @@ -836,10 +921,7 @@ def __get_pydantic_core_schema__( schema = handler(Any) def _serialize(obj: Colormap) -> Any: - if obj.info is not None and obj.info.qualified_name: - # this is a catalog item - return obj.info.qualified_name - return obj.as_dict() + return obj._catalog_name or obj.as_dict() ser = core_schema.plain_serializer_function_ser_schema(_serialize) return core_schema.no_info_after_validator_function( @@ -956,6 +1038,14 @@ def to_pyqtgraph(self) -> pyqtgraph.ColorMap: return _external.to_pyqtgraph(self) +def _rebuild_colormap( + cls: type[Colormap], value: ColormapLike, kwargs: dict[str, Any] +) -> Colormap: + # pickle's two-tuple reduce form passes positional arguments only, and the rest of + # the colormap's state is keyword-only + return cls(value, **kwargs) + + class ColorStop(NamedTuple): """A color stop in a color gradient. diff --git a/tests/test_colormap.py b/tests/test_colormap.py index 4c358bfeb..8c98bc753 100644 --- a/tests/test_colormap.py +++ b/tests/test_colormap.py @@ -1,4 +1,7 @@ +import inspect +import pickle import sys +from copy import copy, deepcopy from functools import partial from typing import Any @@ -146,9 +149,6 @@ def test_with_extremes_preserves_interpolation() -> None: def test_colormap_copy() -> None: """Test Colormap copy.""" - import pickle - from copy import copy, deepcopy - cmap1 = Colormap(["red", "magenta", "blue"]) cmap2 = Colormap(["r", (0.5, "magenta"), "b"]) assert cmap1 == cmap2 @@ -419,6 +419,146 @@ def test_with_extremes() -> None: assert "under" in cm2._repr_html_() +def _configured() -> Colormap: + return Colormap( + ["red", "blue"], + name="foo", + category="sequential", + interpolation="nearest", + under="green", + over="yellow", + bad="black", + neg_inf="cyan", + pos_inf="magenta", + nan="white", + ) + + +def test_with_extremes_preserves_what_is_not_passed() -> None: + cm = Colormap(["red", "blue"], identifier="my_id", under="green", nan="white") + + new = cm.with_extremes(over="yellow") + + assert new.under_color == Color("green") + assert new.nan_color == Color("white") + assert new.over_color == Color("yellow") + assert new.identifier == "my_id" + + +def test_reversed_preserves_state_and_swaps_the_directional_colors() -> None: + cm = _configured() + # the comparisons with `cm` below only catch a dropped or unswapped field if + # every field is non-default and the extreme colors are all distinct + extremes = [cm.under_color, cm.over_color, cm.bad_color] + extremes += [cm.neg_inf_color, cm.pos_inf_color, cm.nan_color] + assert None not in extremes and len(set(extremes)) == len(extremes) + assert cm.category is not None and cm.interpolation != "linear" + + rev = cm.reversed() + + assert rev.name == f"{cm.name}_r" + assert rev.category == cm.category + assert rev.interpolation == cm.interpolation + # under/over and neg_inf/pos_inf name the ends they extend, so they follow them + assert rev.under_color == cm.over_color + assert rev.over_color == cm.under_color + assert rev.neg_inf_color == cm.pos_inf_color + assert rev.pos_inf_color == cm.neg_inf_color + assert rev.bad_color == cm.bad_color + assert rev.nan_color == cm.nan_color + assert rev.reversed() == cm + + +# one catalog entry per kind of state that "_r" and reversed() could disagree on +@pytest.mark.parametrize( + "name", + [ + "napari:HiLo", # under/over in the record, which must swap + "tab10", # nearest interpolation + "napari:nan", # bad in the record + ], +) +def test_reversed_agrees_with_the_r_suffix(name: str) -> None: + assert Colormap(name).reversed() == Colormap(f"{name}_r") + + +def test_r_suffix_keeps_explicit_directional_colors() -> None: + # the record's colors swap; an explicit argument names the end it lands on + assert Colormap("napari:HiLo_r", under="green").under_color == Color("green") + + +def test_pickle_preserves_constructor_state() -> None: + cm = _configured() + + rt = pickle.loads(pickle.dumps(cm)) # noqa: S301 + + assert rt == cm + assert (rt.name, rt.identifier, rt.category, rt.interpolation) == ( + cm.name, + cm.identifier, + cm.category, + cm.interpolation, + ) + + +@pytest.mark.parametrize("name", ["viridis", "viridis_r", "napari:HiLo_r"]) +def test_copies_of_a_catalog_colormap_keep_its_info(name: str) -> None: + cm = Colormap(name) + + for copied in (pickle.loads(pickle.dumps(cm)), copy(cm), deepcopy(cm)): # noqa: S301 + assert copied == cm + assert copied.name == cm.name + assert copied.info is not None and copied.info == cm.info + + +# a sample for every constructor argument besides `value`. A new argument fails the +# test below until it has a sample here, and that sample fails it until the argument +# is in the overrides that `Colormap.__init__` checks before recording a catalog name. +_ARGUMENT_SAMPLES: dict[str, Any] = { + "name": "renamed", + "identifier": "renamed", + "category": "sequential", + "interpolation": False, + "under": "red", + "over": "red", + "bad": "red", + "neg_inf": "red", + "pos_inf": "red", + "nan": "red", + "cmap_kwargs": {"start": 1.0}, +} + + +def test_any_constructor_argument_disables_the_catalog_name() -> None: + params = set(inspect.signature(Colormap).parameters) - {"value"} + assert params == set(_ARGUMENT_SAMPLES) + + assert Colormap("cubehelix")._catalog_name == "cubehelix:cubehelix" + for arg, sample in _ARGUMENT_SAMPLES.items(): + cm = Colormap("cubehelix", **{arg: sample}) + assert cm._catalog_name is None, f"{arg}= is not treated as an override" + + +def test_copies_of_a_callable_colormap_stay_callable() -> None: + # as_dict() samples the function into 256 stops; a copy must keep the function + cm = Colormap("cubehelix", cmap_kwargs={"start": 1.0, "rotation": -1.0}) + + for copied in (pickle.loads(pickle.dumps(cm)), copy(cm), deepcopy(cm)): # noqa: S301 + npt.assert_array_equal(copied.lut(17, gamma=2), cm.lut(17, gamma=2)) + + +def test_as_dict_round_trip() -> None: + plain = Colormap(["red", "blue"]) + assert set(plain.as_dict()) == {"name", "identifier", "category", "value"} + + cm = _configured() + d = cm.as_dict() + + assert d["interpolation"] == "nearest" + assert d["nan"] == list(Color("white")) + assert Colormap(**d) == cm + + def test_shifted() -> None: cm = Colormap(["red", "blue", "yellow"]) assert cm.shifted(1) == cm diff --git a/tests/test_model_fields.py b/tests/test_model_fields.py index d8e1bc694..28e1640c1 100644 --- a/tests/test_model_fields.py +++ b/tests/test_model_fields.py @@ -58,3 +58,42 @@ class MyModel(psygnal.EventedModel): # type: ignore color=np.array([1, 0, 0]), colormap=["r", (0.7, "b"), "w"], stops="green_r" ) assert MyModel.model_validate_json(obj.model_dump_json()) == obj + + +def test_pydantic_preserves_modified_catalog_colormaps() -> None: + class MyModel(pydantic.BaseModel): + colormap: Colormap + + def round_trip(cmap: Colormap) -> Colormap: + obj = MyModel(colormap=cmap) + return MyModel.model_validate_json(obj.model_dump_json()).colormap + + # an unmodified catalog colormap still serializes to its qualified name alone + obj = MyModel(colormap=Colormap("viridis")) + assert '"colormap":"bids:viridis"' in obj.model_dump_json() + + assert round_trip(Colormap("viridis", under="red")).under_color == Color("red") + assert round_trip(Colormap("viridis", name="renamed")).name == "renamed" + assert round_trip(Colormap("viridis_r")) == Colormap("viridis_r") + # ... and the reversed name alone rebuilds it, so it serializes to that + obj = MyModel(colormap=Colormap("viridis_r")) + assert '"colormap":"bids:viridis_r"' in obj.model_dump_json() + # a falsy argument is still an override: False is "nearest", viridis is linear + nearest = Colormap("viridis", interpolation=False) + assert round_trip(nearest).interpolation == "nearest" + + # a parametrization too small for the tolerant __eq__ to see still must not + # collapse to the plain catalog name + tweaked = Colormap("cubehelix", cmap_kwargs={"start": 0.500001}) + assert round_trip(tweaked).as_dict() == tweaked.as_dict() + + +def test_psygnal_serialization_of_a_configured_colormap() -> None: + psygnal = pytest.importorskip("psygnal") + + class MyModel(psygnal.EventedModel): # type: ignore + colormap: Colormap + + cmap = Colormap(["r", "b"], interpolation="nearest", under="green", nan="orange") + obj = MyModel(colormap=cmap) + assert MyModel.model_validate_json(obj.model_dump_json()).colormap == cmap