From 0915f0f6e57b3de51d31e6c71c9c165d4f0bbf71 Mon Sep 17 00:00:00 2001 From: Alister Burt Date: Wed, 14 Jan 2026 19:12:40 -0800 Subject: [PATCH 1/4] add parametrized colormap registry and expose cubehelix with Colormap("cubehelix", cmap_kwargs: dict[str, Any]) API --- src/cmap/_colormap.py | 17 +++++++++++- src/cmap/_parametrized_colormaps.py | 42 +++++++++++++++++++++++++++++ tests/test_colormap.py | 15 +++++++++++ 3 files changed, 73 insertions(+), 1 deletion(-) create mode 100644 src/cmap/_parametrized_colormaps.py diff --git a/src/cmap/_colormap.py b/src/cmap/_colormap.py index 748742ddc..741c451ef 100644 --- a/src/cmap/_colormap.py +++ b/src/cmap/_colormap.py @@ -15,6 +15,7 @@ from . import _external from ._catalog import Catalog from ._color import Color, ColorLike +from ._parametrized_colormaps import get_parametrized_colormap_function if TYPE_CHECKING: from collections.abc import Iterator @@ -92,6 +93,8 @@ class Colormap: - a `str` containing a recognized string colormap name (e.g. `"viridis"`, `"magma"`), optionally suffixed with `"_r"` to reverse the colormap (e.g. `"viridis"`, `"magma_r"`). + - a `str` containing a registered colormap function name (e.g. `"cubehelix"`), + used with the `cmap_kwargs` parameter to pass function arguments. - An iterable of [ColorLike](../../colors.md#colorlike-objects) values (any object that can be cast to a [`Color`][cmap.Color]), or "color-stop-like" tuples ( `(float, ColorLike)` where the first element is a scalar value @@ -109,6 +112,12 @@ class Colormap: the matplotlib docs for more. - a `Callable` that takes an array of N values in the range [0, 1] and returns an (N, 4) array of RGBA values in the range [0, 1]. + cmap_kwargs : dict[str, Any] | None + Keyword arguments to pass to a colormap function when `value` is a string + naming a registered function (e.g. `"cubehelix"`). For example: + `Colormap("cubehelix", cmap_kwargs={"start": 1.0, "rotation": -1.0})`. + If provided when `value` is not a registered function name, a `ValueError` + will be raised. name : str | None A name for the colormap. If None, will be set to the identifier or the string `"custom colormap"`. @@ -226,10 +235,16 @@ def __init__( under: ColorLike | None = None, over: ColorLike | None = None, bad: ColorLike | None = None, + cmap_kwargs: dict[str, Any] | None = None, ) -> None: self.info: CatalogItem | None = None - if isinstance(value, str): + if isinstance(value, str) and cmap_kwargs is not None: + name = name or value + colormap_func = get_parametrized_colormap_function(name) + colormap_func = partial(colormap_func, **cmap_kwargs) + stops = _parse_colorstops(colormap_func) + elif isinstance(value, str): rev = value.endswith("_r") info = self.catalog()[value[:-2] if rev else value] name = name or f"{info.namespace}:{info.name}" diff --git a/src/cmap/_parametrized_colormaps.py b/src/cmap/_parametrized_colormaps.py new file mode 100644 index 000000000..eee825616 --- /dev/null +++ b/src/cmap/_parametrized_colormaps.py @@ -0,0 +1,42 @@ +"""Registry of parametrized colormap functions.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable + +from .data.cubehelix import cubehelix + +if TYPE_CHECKING: + from typing import Any + + import numpy as np + + LutCallable = Callable[[np.ndarray], np.ndarray] + +# Registry of colormap functions that accept parameters +FUNCTION_REGISTRY: dict[str, Callable[..., Any]] = { + "cubehelix": cubehelix, +} + + +def get_parametrized_colormap_function(name: str) -> Callable[..., Any]: + """ + Get a parametrized colormap function by name. + + Args: + name: Name of the colormap function + + Returns + ------- + The colormap function + + Raises + ------ + ValueError: If the function name is not in the registry + """ + if name not in FUNCTION_REGISTRY: + available = ", ".join(sorted(FUNCTION_REGISTRY.keys())) + raise ValueError( + f"Unknown colormap function: {name!r}. Available functions: {available}" + ) + return FUNCTION_REGISTRY[name] diff --git a/tests/test_colormap.py b/tests/test_colormap.py index 3a9484d71..bd63c428a 100644 --- a/tests/test_colormap.py +++ b/tests/test_colormap.py @@ -236,3 +236,18 @@ def test_shifted() -> None: assert cm.shifted(0.5).shifted(-0.5) == cm # two shifts of 0.5 should give the original array assert cm.shifted().shifted() == cm + + +def test_function_colormap_with_cmap_kwargs() -> None: + # construct cubehelix with custom parameters + ch = Colormap("cubehelix", cmap_kwargs={"start": 1.0, "rotation": -1.0}) + + # values should be different from default cubehelix + default_ch = Colormap("cubehelix") + assert ch(0.5) != default_ch(0.5) + + +def test_invalid_function_colormap_with_cmap_kwargs() -> None: + # name which doesn't map to a registered function should raise ValueError + with pytest.raises(ValueError, match="Unknown colormap function"): + Colormap("nonexistent_function", cmap_kwargs={"param": 1.0}) From ec4c8918aa14eec4b83a6b638187809f92be1d83 Mon Sep 17 00:00:00 2001 From: Talley Lambert Date: Tue, 3 Feb 2026 20:52:15 -0500 Subject: [PATCH 2/4] remove registry --- src/cmap/_colormap.py | 23 ++++++++++------ src/cmap/_parametrized_colormaps.py | 42 ----------------------------- tests/test_colormap.py | 6 ++--- 3 files changed, 18 insertions(+), 53 deletions(-) delete mode 100644 src/cmap/_parametrized_colormaps.py diff --git a/src/cmap/_colormap.py b/src/cmap/_colormap.py index 741c451ef..2e6755281 100644 --- a/src/cmap/_colormap.py +++ b/src/cmap/_colormap.py @@ -15,7 +15,6 @@ from . import _external from ._catalog import Catalog from ._color import Color, ColorLike -from ._parametrized_colormaps import get_parametrized_colormap_function if TYPE_CHECKING: from collections.abc import Iterator @@ -239,12 +238,7 @@ def __init__( ) -> None: self.info: CatalogItem | None = None - if isinstance(value, str) and cmap_kwargs is not None: - name = name or value - colormap_func = get_parametrized_colormap_function(name) - colormap_func = partial(colormap_func, **cmap_kwargs) - stops = _parse_colorstops(colormap_func) - elif isinstance(value, str): + if isinstance(value, str): rev = value.endswith("_r") info = self.catalog()[value[:-2] if rev else value] name = name or f"{info.namespace}:{info.name}" @@ -253,6 +247,14 @@ def __init__( under = info.under if under is None else under bad = info.bad if bad is None else bad self.info = info + + # Check if cmap_kwargs is provided for a non-callable colormap + if cmap_kwargs and not callable(info.data): + raise TypeError( + f"Cannot apply cmap_kwargs to colormap {info.name!r}: " + "colormap is not a parametrized callable" + ) + if isinstance(info.data, list): if not info.data: # pragma: no cover raise ValueError(f"Catalog colormap {info.name!r} has no data") @@ -267,7 +269,12 @@ def __init__( f"Invalid catalog colormap data for {info.name!r}: {info.data}" ) else: - stops = _parse_colorstops(info.data) + _data: ColormapLike + if cmap_kwargs: + _data = partial(cast("Callable", info.data), **cmap_kwargs) + else: + _data = info.data + stops = _parse_colorstops(_data) if interpolation is None: interpolation = info.interpolation if rev: diff --git a/src/cmap/_parametrized_colormaps.py b/src/cmap/_parametrized_colormaps.py deleted file mode 100644 index eee825616..000000000 --- a/src/cmap/_parametrized_colormaps.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Registry of parametrized colormap functions.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Callable - -from .data.cubehelix import cubehelix - -if TYPE_CHECKING: - from typing import Any - - import numpy as np - - LutCallable = Callable[[np.ndarray], np.ndarray] - -# Registry of colormap functions that accept parameters -FUNCTION_REGISTRY: dict[str, Callable[..., Any]] = { - "cubehelix": cubehelix, -} - - -def get_parametrized_colormap_function(name: str) -> Callable[..., Any]: - """ - Get a parametrized colormap function by name. - - Args: - name: Name of the colormap function - - Returns - ------- - The colormap function - - Raises - ------ - ValueError: If the function name is not in the registry - """ - if name not in FUNCTION_REGISTRY: - available = ", ".join(sorted(FUNCTION_REGISTRY.keys())) - raise ValueError( - f"Unknown colormap function: {name!r}. Available functions: {available}" - ) - return FUNCTION_REGISTRY[name] diff --git a/tests/test_colormap.py b/tests/test_colormap.py index bd63c428a..7e7e45b52 100644 --- a/tests/test_colormap.py +++ b/tests/test_colormap.py @@ -248,6 +248,6 @@ def test_function_colormap_with_cmap_kwargs() -> None: def test_invalid_function_colormap_with_cmap_kwargs() -> None: - # name which doesn't map to a registered function should raise ValueError - with pytest.raises(ValueError, match="Unknown colormap function"): - Colormap("nonexistent_function", cmap_kwargs={"param": 1.0}) + # providing cmap_kwargs to a non-callable colormap should raise TypeError + with pytest.raises(TypeError, match=r"Cannot apply cmap_kwargs.*not callable"): + Colormap("viridis", cmap_kwargs={"param": 1.0}) From 437ba117c2d91a7461439b402c941901fa9e9e0d Mon Sep 17 00:00:00 2001 From: Talley Lambert Date: Tue, 3 Feb 2026 20:56:39 -0500 Subject: [PATCH 3/4] update docs --- src/cmap/_colormap.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/cmap/_colormap.py b/src/cmap/_colormap.py index 2e6755281..433072c2f 100644 --- a/src/cmap/_colormap.py +++ b/src/cmap/_colormap.py @@ -92,8 +92,6 @@ class Colormap: - a `str` containing a recognized string colormap name (e.g. `"viridis"`, `"magma"`), optionally suffixed with `"_r"` to reverse the colormap (e.g. `"viridis"`, `"magma_r"`). - - a `str` containing a registered colormap function name (e.g. `"cubehelix"`), - used with the `cmap_kwargs` parameter to pass function arguments. - An iterable of [ColorLike](../../colors.md#colorlike-objects) values (any object that can be cast to a [`Color`][cmap.Color]), or "color-stop-like" tuples ( `(float, ColorLike)` where the first element is a scalar value @@ -112,11 +110,11 @@ class Colormap: - a `Callable` that takes an array of N values in the range [0, 1] and returns an (N, 4) array of RGBA values in the range [0, 1]. cmap_kwargs : dict[str, Any] | None - Keyword arguments to pass to a colormap function when `value` is a string - naming a registered function (e.g. `"cubehelix"`). For example: + Keyword arguments to pass to the colormap function when `value` is a + registered colormap name that maps to a callable function. For example: `Colormap("cubehelix", cmap_kwargs={"start": 1.0, "rotation": -1.0})`. - If provided when `value` is not a registered function name, a `ValueError` - will be raised. + If provided when `value` does not resolve to a callable colormap function, a + `TypeError` will be raised. name : str | None A name for the colormap. If None, will be set to the identifier or the string `"custom colormap"`. @@ -137,6 +135,14 @@ class Colormap: The color to use for values above the colormap's range. bad : ColorLike | None The color to use for bad (NaN, inf) values. + + Raises + ------ + TypeError + If `cmap_kwargs` is provided and `value` does not resolve to a callable + colormap function. + ValueError + If `value` is a string colormap name that is not found in the catalog. """ __slots__ = ( From 7e0f94f7c662b3a3d0b58679427ac96570dc6f30 Mon Sep 17 00:00:00 2001 From: Talley Lambert Date: Tue, 3 Feb 2026 20:58:18 -0500 Subject: [PATCH 4/4] fix test --- tests/test_colormap.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_colormap.py b/tests/test_colormap.py index 7e7e45b52..d1550136b 100644 --- a/tests/test_colormap.py +++ b/tests/test_colormap.py @@ -249,5 +249,5 @@ def test_function_colormap_with_cmap_kwargs() -> None: def test_invalid_function_colormap_with_cmap_kwargs() -> None: # providing cmap_kwargs to a non-callable colormap should raise TypeError - with pytest.raises(TypeError, match=r"Cannot apply cmap_kwargs.*not callable"): + with pytest.raises(TypeError, match=r"Cannot apply cmap_kwargs to colormap"): Colormap("viridis", cmap_kwargs={"param": 1.0})