From 9380611161ac87b56f05954acc3c40463272e3b3 Mon Sep 17 00:00:00 2001 From: Matthias Schabel Date: Mon, 10 Aug 2026 13:27:46 -0700 Subject: [PATCH 1/2] feat: add per-class colors for infinities, NaN, and masked entries cmap colors three exceptional classes: under, over, and bad. Floating point data has more. Negative and positive infinity are indistinguishable from ordinary out-of-range values, and NaN is indistinguishable from a masked entry. Adds neg_inf, pos_inf, nan, and masked. Each falls back to the color its class uses now: neg_inf to under, pos_inf to over, nan and masked to bad. bad is kept as the joint fallback for both of its children, so code that sets it is unaffected and either child may be set alone. Routing appends four fallback-resolved rows to a call-local copy of the over/under LUT, so a class with no color of its own lands on exactly the row it lands on now. Colormap.lut() is unchanged. The infinity masks are taken before the input is scaled by N: that multiply overflows large finite values to infinity (float16 65504 does it), and those are out of range rather than infinite. Co-Authored-By: Claude Opus 5 (1M context) Reviewed-By: Codex (gpt-5.6-sol, reasoning effort xhigh) --- src/cmap/_colormap.py | 145 ++++++++++++++++++++++++++++++++++---- src/cmap/_external.py | 6 +- tests/test_colormap.py | 98 ++++++++++++++++++++++++++ tests/test_third_party.py | 4 ++ 4 files changed, 238 insertions(+), 15 deletions(-) diff --git a/src/cmap/_colormap.py b/src/cmap/_colormap.py index 599e18ef9..b576e9613 100644 --- a/src/cmap/_colormap.py +++ b/src/cmap/_colormap.py @@ -141,6 +141,18 @@ class Colormap: The color to use for NaN and masked values. When no bad color is set, they are transparent. Note that infinities are not bad values here: they use `under` and `over`. + neg_inf : ColorLike | None + The color to use for negative infinity. When unset, negative infinity uses + `under`. + pos_inf : ColorLike | None + The color to use for positive infinity. When unset, positive infinity uses + `over`. + nan : ColorLike | None + The color to use for NaN. When unset, NaN uses `bad`. + masked : ColorLike | None + The color to use for entries masked by a `numpy.ma` masked array. When unset, + masked entries use `bad`. A masked entry takes this color whatever value it + hides, so a masked infinity is masked rather than infinite. Raises ------ @@ -153,6 +165,7 @@ class Colormap: __slots__ = ( "__weakref__", + "_has_exceptional", "_initialized", "_lut_cache", "bad_color", @@ -161,8 +174,12 @@ class Colormap: "identifier", "info", "interpolation", + "masked_color", "name", + "nan_color", + "neg_inf_color", "over_color", + "pos_inf_color", "under_color", ) @@ -224,6 +241,24 @@ class Colormap: If provided, and `Colormap.lut` is called with `with_over_under=True`, `bad_color` will be the last color in the LUT (`lut[-1]`). + + `nan_color` and `masked_color` override it for their own class. It remains the + color both of them fall back to. + """ + + neg_inf_color: Color | None + """A color to use for negative infinity, overriding `under_color`.""" + + pos_inf_color: Color | None + """A color to use for positive infinity, overriding `over_color`.""" + + nan_color: Color | None + """A color to use for NaN, overriding `bad_color`.""" + + masked_color: Color | None + """A color to use for masked entries, overriding `bad_color`. + + Applies to any entry masked by a `numpy.ma` masked array, whatever value it hides. """ _catalog_instance: Catalog | None = None @@ -246,6 +281,10 @@ def __init__( under: ColorLike | None = None, over: ColorLike | None = None, bad: ColorLike | None = None, + neg_inf: ColorLike | None = None, + pos_inf: ColorLike | None = None, + nan: ColorLike | None = None, + masked: ColorLike | None = None, cmap_kwargs: dict[str, Any] | None = None, ) -> None: self.info: CatalogItem | None = None @@ -318,6 +357,20 @@ def __init__( self.under_color = None if under is None else Color(under) self.over_color = None if over is None else Color(over) self.bad_color = None if bad is None else Color(bad) + self.neg_inf_color = None if neg_inf is None else Color(neg_inf) + self.pos_inf_color = None if pos_inf is None else Color(pos_inf) + self.nan_color = None if nan is None else Color(nan) + self.masked_color = None if masked is None else Color(masked) + # a colormap with none of these takes the same path it did before they existed + self._has_exceptional = any( + c is not None + for c in ( + self.neg_inf_color, + self.pos_inf_color, + self.nan_color, + self.masked_color, + ) + ) self._lut_cache: dict[LutCacheKey, np.ndarray] = {} self._initialized = True @@ -361,12 +414,14 @@ def __call__( For float input, values outside the [0, 1] range and values that are not finite do not map into the ramp: - - values below 0, and negative infinity, use `under_color` (when unset, the - first color in the colormap). - - values above 1, and positive infinity, use `over_color` (when unset, the - last color in the colormap). - - NaN, and entries masked by a `numpy.ma` masked array, use `bad_color` - (when unset, transparent). + - values below 0 use `under_color` (when unset, the first color in the + colormap), and values above 1 use `over_color` (when unset, the last). + - negative and positive infinity use `neg_inf_color` and `pos_inf_color` + (when unset, `under_color` and `over_color`). + - NaN uses `nan_color`, and entries masked by a `numpy.ma` masked array use + `masked_color` (when either is unset, `bad_color`, which is itself + transparent when unset). A masked entry takes the masked color whatever + value it hides. For integer input, which indexes the LUT directly, an index at or beyond N uses `over_color`, and a negative index uses `under_color` rather than @@ -410,16 +465,24 @@ def __call__( >>> colored_img = cmap(data) """ lut = self.lut(N=N, gamma=gamma, with_over_under=True) - if bytes: - lut = (lut * 255).astype(np.uint8) # the lut will have three additional colors at the end for under, over, and bad N = len(lut) - 3 + if self._has_exceptional: + lut = self._with_exceptional_colors(lut) + if bytes: + lut = (lut * 255).astype(np.uint8) xa = np.array(x, copy=True) if not xa.dtype.isnative: # Native byteorder is faster. xa = xa.byteswap().view(xa.dtype.newbyteorder()) - if xa.dtype.kind == "f": + is_float = xa.dtype.kind == "f" + if self._has_exceptional and is_float: + # before the scaling below: it overflows large finite values to infinity + # (float16 65504), and those are out of range rather than infinite. + mask_neg_inf = np.isneginf(xa) + mask_pos_inf = np.isposinf(xa) + if is_float: xa *= N # xa == 1 (== N after multiplication) is not out of range. xa[xa == N] = N - 1 @@ -429,11 +492,12 @@ def __call__( # If input was masked, start from its mask: a masked array can still carry # unmasked nans. `|` rather than `|=`, so x's own mask isn't written to. if np.ma.is_masked(x): - mask_bad = x.mask # type: ignore - if xa.dtype.kind == "f": - mask_bad = mask_bad | np.isnan(xa) + mask_masked = x.mask # type: ignore + mask_nan = np.isnan(xa) if is_float else False + mask_bad = (mask_masked | mask_nan) if is_float else mask_masked else: - mask_bad = np.isnan(xa) + mask_masked = False + mask_nan = mask_bad = np.isnan(xa) with np.errstate(invalid="ignore"): # We need this cast for unsigned ints as well as floats @@ -442,16 +506,45 @@ def __call__( xa[mask_under] = N xa[mask_over] = N + 1 xa[mask_bad] = N + 2 + if self._has_exceptional: + # last wins: a masked entry is masked whatever value it hides + if is_float: + xa[mask_neg_inf] = N + 3 + xa[mask_pos_inf] = N + 4 + xa[mask_nan] = N + 5 + xa[mask_masked] = N + 6 rgba = lut.take(xa, axis=0, mode="clip") return rgba if np.iterable(x) else Color(rgba) + def _with_exceptional_colors(self, lut: np.ndarray) -> np.ndarray: + """Return `lut` with four rows appended, one per exceptional value class. + + Each appended row falls back to the row its class would otherwise have used, + so routing a class to its own row cannot change any color while that class + has no color of its own. `lut` must be an over/under LUT. + """ + under, over, bad = lut[-3], lut[-2], lut[-1] + return np.vstack( + ( + lut, + under if self.neg_inf_color is None else self.neg_inf_color.rgba, + over if self.pos_inf_color is None else self.pos_inf_color.rgba, + bad if self.nan_color is None else self.nan_color.rgba, + bad if self.masked_color is None else self.masked_color.rgba, + ) + ) + def with_extremes( self, *, bad: ColorLike | None = None, under: ColorLike | None = None, over: ColorLike | None = None, + neg_inf: ColorLike | None = None, + pos_inf: ColorLike | None = None, + nan: ColorLike | None = None, + masked: ColorLike | None = None, ) -> Colormap: """Return a copy of the colormap with new extreme values.""" return type(self)( @@ -462,6 +555,10 @@ def with_extremes( bad=bad, under=under, over=over, + neg_inf=neg_inf, + pos_inf=pos_inf, + nan=nan, + masked=masked, ) def as_dict(self) -> ColormapDict: @@ -609,6 +706,10 @@ def shifted( under=self.under_color, over=self.over_color, bad=self.bad_color, + neg_inf=self.neg_inf_color, + pos_inf=self.pos_inf_color, + nan=self.nan_color, + masked=self.masked_color, ) def to_css( @@ -670,6 +771,10 @@ def __eq__(self, other: object) -> bool: and self.under_color == other.under_color and self.over_color == other.over_color and self.bad_color == other.bad_color + and self.neg_inf_color == other.neg_inf_color + and self.pos_inf_color == other.pos_inf_color + and self.nan_color == other.nan_color + and self.masked_color == other.masked_color and self.interpolation == other.interpolation ) @@ -718,6 +823,20 @@ def _repr_html_(self) -> str: f"over {_html_color_patch(self.over_color)}" "" ) + if self._has_exceptional: + patches = ( + ("neg_inf", self.neg_inf_color), + ("pos_inf", self.pos_inf_color), + ("nan", self.nan_color), + ("masked", self.masked_color), + ) + swatches = " ".join( + f"{name} {_html_color_patch(c)}" for name, c in patches if c is not None + ) + html += ( + '
' + f"{swatches}
" + ) return html diff --git a/src/cmap/_external.py b/src/cmap/_external.py index 68a5bd8ce..23ed001a4 100644 --- a/src/cmap/_external.py +++ b/src/cmap/_external.py @@ -109,8 +109,10 @@ def to_napari(cm: Colormap) -> NapariColormap: kwargs["interpolation"] = ( "zero" if cm.interpolation == "nearest" else "linear" ) - if "nan_color" in param_names and cm.bad_color is not None: - kwargs["nan_color"] = cm.bad_color.rgba + # napari's nan_color covers nan alone, so prefer cmap's nan color over bad + nan_color = cm.nan_color or cm.bad_color + if "nan_color" in param_names and nan_color is not None: + kwargs["nan_color"] = nan_color.rgba if "high_color" in param_names and cm.over_color is not None: kwargs["high_color"] = cm.over_color.rgba if "low_color" in param_names and cm.under_color is not None: diff --git a/tests/test_colormap.py b/tests/test_colormap.py index 41524269e..bad939546 100644 --- a/tests/test_colormap.py +++ b/tests/test_colormap.py @@ -187,6 +187,104 @@ def test_colormap_masked_array_with_unmasked_nan() -> None: npt.assert_array_equal(cmap(all_false), cmap(np.array([0.25, np.nan]))) +def test_exceptional_colors() -> None: + cmap = Colormap( + ["red", "blue"], + under="green", + over="yellow", + bad="black", + neg_inf="cyan", + pos_inf="magenta", + nan="white", + masked="orange", + ) + mask = [False] * 6 + [True] * 3 + data = np.ma.masked_array( + [-np.inf, -0.5, 0.5, 1.5, np.inf, np.nan, np.inf, np.nan, 0.25], mask=mask + ) + expect = np.array( + [ + Color("cyan").rgba, # -inf + Color("green").rgba, # under range, finite + Colormap(["red", "blue"])(0.5).rgba, # in range + Color("yellow").rgba, # over range, finite + Color("magenta").rgba, # +inf + Color("white").rgba, # nan, not masked + Color("orange").rgba, # masked wins over +inf + Color("orange").rgba, # masked wins over nan + Color("orange").rgba, # masked + ] + ) + + npt.assert_array_equal(cmap(data), expect) + npt.assert_array_equal(cmap(data, bytes=True), (expect * 255).astype(np.uint8)) + npt.assert_array_equal(data.mask, mask) + + +def test_exceptional_colors_fall_back_to_the_legacy_extremes() -> None: + cmap = Colormap(["red", "blue"], under="green", over="yellow", bad="black") + data = np.ma.masked_array( + [-np.inf, np.inf, np.nan, 0.5], mask=[False, False, False, True] + ) + legacy = [ + Color("green").rgba, # -inf -> under + Color("yellow").rgba, # +inf -> over + Color("black").rgba, # nan -> bad + Color("black").rgba, # masked -> bad + ] + npt.assert_array_equal(cmap(data), legacy) + + # setting one leaves the other three on their legacy destinations + one = cmap.with_extremes(under="green", over="yellow", bad="black", nan="white") + npt.assert_array_equal(one(data), [*legacy[:2], Color("white").rgba, legacy[3]]) + + +def test_finite_values_that_overflow_when_scaled_are_not_infinite() -> None: + # float16 65504 becomes inf once multiplied by N, but it is over-range, not infinite + cmap = Colormap( + ["red", "blue"], under="green", over="yellow", neg_inf="cyan", pos_inf="magenta" + ) + data = np.array([65504, -65504, np.inf, -np.inf], dtype=np.float16) + with np.errstate(over="ignore"): + rgba = cmap(data) + npt.assert_array_equal( + rgba, + [ + Color("yellow").rgba, + Color("green").rgba, + Color("magenta").rgba, + Color("cyan").rgba, + ], + ) + + +def test_masked_dtypes_keep_their_existing_behavior() -> None: + cmap = Colormap(["red", "blue"], bad="black") + bad = Color("black").rgba + + npt.assert_array_equal(cmap(np.ma.masked_array([0, 1], mask=[True, False]))[0], bad) + + obj = np.ma.masked_array(np.array([0.25, 0.5], dtype=object), mask=[True, False]) + npt.assert_array_equal(cmap(obj)[0], bad) + + # object dtype without a mask reaches np.isnan, which has never accepted it + with pytest.raises(TypeError): + cmap(np.ma.masked_array(np.array([0.25], dtype=object), mask=np.ma.nomask)) + + npt.assert_array_equal(cmap(np.ma.masked_array(0.5, mask=True)).rgba, bad) + + +@pytest.mark.parametrize("field", ["neg_inf", "pos_inf", "nan", "masked"]) +def test_exceptional_colors_are_colormap_state(field: str) -> None: + plain = Colormap(["red", "blue"]) + cmap = Colormap(["red", "blue"], **{field: "orange"}) + + assert cmap != plain + assert plain.with_extremes(**{field: "orange"}) == cmap + assert cmap.shifted(1) == cmap + assert field in cmap._repr_html_() + + def test_fill_stops() -> None: assert _fill_stops([None, None, None]) == [0, 0.5, 1.0] assert _fill_stops([None, 0.8, None]) == [0, 0.8, 1.0] diff --git a/tests/test_third_party.py b/tests/test_third_party.py index a161dac75..f02cbc4b3 100644 --- a/tests/test_third_party.py +++ b/tests/test_third_party.py @@ -94,6 +94,10 @@ def test_napari_extreme_colors() -> None: np.testing.assert_allclose(ncm.high_color, Color(OVER).rgba) np.testing.assert_allclose(ncm.nan_color, Color(bad).rgba) + # a more specific nan color takes precedence over bad + specific = Colormap(["black", "white"], bad=bad, nan="magenta").to_napari() + np.testing.assert_allclose(specific.nan_color, Color("magenta").rgba) + @pytest.mark.skipif( sys.platform == "darwin" and sys.version_info >= (3, 13), From 430248162476c7d02e2160572e35796c312101f8 Mon Sep 17 00:00:00 2001 From: Matthias Schabel Date: Thu, 13 Aug 2026 10:03:36 -0700 Subject: [PATCH 2/2] fix: preserve colormap state through reversed(), with_extremes(), pickle, and as_dict() None of the four channels carried the extreme colors, and reversed() and with_extremes() dropped the interpolation mode as well. Following matplotlib where it has a position: reversed() swaps the directional colors and keeps the rest, with_extremes() keeps anything not passed, __reduce__ carries the constructor state alongside the stops, and as_dict() gains optional keys that are written only when set, so existing payloads are unchanged. The pydantic serializer emitted a catalog colormap's qualified name unconditionally, which discarded any added extremes and turned "viridis_r" back into plain viridis. It now emits the name only when the name alone rebuilds the same colormap. Colormap("x_r") swaps the catalog record's under and over, so it agrees with Colormap("x").reversed(). Co-Authored-By: Claude Opus 5 (1M context) Reviewed-By: Codex (gpt-5.6-sol, reasoning effort xhigh) --- src/cmap/_colormap.py | 123 +++++++++++++++++++++++++++++++------ tests/test_colormap.py | 91 +++++++++++++++++++++++++++ tests/test_model_fields.py | 45 ++++++++++++++ 3 files changed, 240 insertions(+), 19 deletions(-) diff --git a/src/cmap/_colormap.py b/src/cmap/_colormap.py index b576e9613..10e31600a 100644 --- a/src/cmap/_colormap.py +++ b/src/cmap/_colormap.py @@ -44,12 +44,28 @@ class MPLSegmentData(TypedDict, total=False): blue: list[tuple[float, float, float]] | Callable[[np.ndarray], np.ndarray] alpha: list[tuple[float, float, float]] | Callable[[np.ndarray], np.ndarray] - class ColormapDict(TypedDict): + class _RequiredColormapDict(TypedDict): name: str identifier: str category: str | None value: list[tuple[float, list[float]]] + class ColormapDict(_RequiredColormapDict, total=False): + # 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: str + under: list[float] + over: list[float] + bad: list[float] + neg_inf: list[float] + pos_inf: list[float] + nan: list[float] + masked: list[float] + + +# the extreme colors, in constructor-argument spelling; `_color` is the property. +# Copy and serialization paths iterate this rather than repeating the list. +_EXTREME_FIELDS = ("under", "over", "bad", "neg_inf", "pos_inf", "nan", "masked") 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 @@ -294,8 +310,14 @@ def __init__( info = self.catalog()[value[:-2] if rev else value] name = name or f"{info.namespace}:{info.name}" 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 @@ -546,33 +568,65 @@ def with_extremes( nan: ColorLike | None = None, masked: 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. To clear one, construct a new `Colormap`. + """ 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, - masked=masked, + 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, + masked=self.masked_color if masked is None else masked, ) + @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 @@ -650,6 +704,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 @@ -661,8 +718,20 @@ 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, + masked=self.masked_color, ) def shifted( @@ -756,8 +825,13 @@ 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. 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): @@ -854,10 +928,13 @@ 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() + state = obj.as_dict() + if obj.info is not None and (qualified := obj.info.qualified_name): + # the name alone is a complete serialization only when it rebuilds the + # same colormap; "viridis_r" and a modified viridis both fail that + if state == Colormap(qualified).as_dict(): + return qualified + return state ser = core_schema.plain_serializer_function_ser_schema(_serialize) return core_schema.no_info_after_validator_function( @@ -974,6 +1051,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 bad939546..31213a521 100644 --- a/tests/test_colormap.py +++ b/tests/test_colormap.py @@ -369,6 +369,97 @@ 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", + masked="orange", + ) + + +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() + + rev = cm.reversed() + + assert rev.name == "foo_r" + assert rev.category == "sequential" + assert rev.interpolation == "nearest" + # under/over and neg_inf/pos_inf name the ends they extend, so they follow them + assert (rev.under_color, rev.over_color) == (Color("yellow"), Color("green")) + assert (rev.neg_inf_color, rev.pos_inf_color) == (Color("magenta"), Color("cyan")) + assert (rev.bad_color, rev.nan_color, rev.masked_color) == ( + Color("black"), + Color("white"), + Color("orange"), + ) + assert rev.reversed() == cm + + +def test_reversed_agrees_with_the_r_suffix() -> None: + assert Colormap("napari:HiLo").reversed() == Colormap("napari:HiLo_r") + # 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: + import pickle + + 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, + ) + + +def test_copies_of_a_callable_colormap_stay_callable() -> None: + import pickle + from copy import copy, deepcopy + + # 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["masked"] == list(Color("orange")) + 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 43c42add1..bdee00dab 100644 --- a/tests/test_model_fields.py +++ b/tests/test_model_fields.py @@ -85,3 +85,48 @@ class MyModel(psygnal.EventedModel): # type: ignore assert MyModel.model_validate_json(data) == obj else: assert MyModel.parse_raw(data) == obj + + +def test_pydantic_preserves_modified_catalog_colormaps() -> None: + class MyModel(BaseModel): + colormap: Colormap + + def round_trip(cmap: Colormap) -> Colormap: + obj = MyModel(colormap=cmap) + data = obj.model_dump_json() if V2 else obj.json() + model = ( + MyModel.model_validate_json(data) + if hasattr(MyModel, "model_validate_json") + else MyModel.parse_raw(data) + ) + return model.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() if V2 else obj.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") + + # 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", masked="orange") + obj = MyModel(colormap=cmap) + + data = obj.model_dump_json() if V2 else obj.json() + + if hasattr(MyModel, "model_validate_json"): + assert MyModel.model_validate_json(data).colormap == cmap + else: + assert MyModel.parse_raw(data).colormap == cmap