diff --git a/ultraplot/colors.py b/ultraplot/colors.py index 0d3de4102..4a48edded 100644 --- a/ultraplot/colors.py +++ b/ultraplot/colors.py @@ -41,6 +41,11 @@ ) from .utils import set_alpha, to_hex, to_rgb, to_rgba, to_xyz, to_xyza +try: + from typing import override +except: + from typing_extensions import override + __all__ = [ "DiscreteColormap", "ContinuousColormap", @@ -3090,13 +3095,10 @@ def __init__(self, kwargs): kwargs : dict-like The source dictionary. """ - super().__init__(kwargs) - # The colormap is initialized with all the base colormaps - # We have to change the classes internally to Perceptual, Continuous or Discrete - # such that ultraplot knows what these objects are. We piggy back on the registering mechanism - # by overriding matplotlib's behavior - for name in tuple(self._cmaps.keys()): - self.register(self._cmaps[name], name=name) + super().__init__({k.lower(): v for k, v in kwargs.items()}) + # The colormap is initialized with all the base colormaps. + # These are converted to ultraplot's own colormap objects + # on the fly when they are first accessed. def _translate_deprecated(self, key): """ @@ -3171,6 +3173,34 @@ def _translate_key(self, original_key, mirror=True): def _has_item(self, key): return key in self._cmaps + def _load_and_register_cmap(self, key, value): + """ + Load a colormap from a file and register it. + """ + path = value["path"] + type = value["type"] + is_default = value.get("is_default", False) + if type == "continuous": + cmap = ContinuousColormap.from_file(path, warn_on_failure=True) + elif type == "discrete": + cmap = DiscreteColormap.from_file(path, warn_on_failure=True) + else: + raise ValueError( + f"Invalid colormap type {type!r} for key {key!r} in file {path!r}. " + "Expected 'continuous' or 'discrete'." + ) + + if cmap: + if is_default and cmap.name.lower() in CMAPS_CYCLIC: + cmap.set_cyclic(True) + self.register(cmap, name=key) + return self._cmaps[key] + else: # failed to load + # remove from registry to avoid trying again + del self._cmaps[key] + warnings._warn_ultraplot(f"Failed to load colormap {key!r} from {path!r}") + return None + def get_cmap(self, cmap): return self.__getitem__(cmap) @@ -3188,9 +3218,23 @@ def __getitem__(self, key): if reverse: key = key.removesuffix("_r") + # Retrieve colormap if self._has_item(key): - value = self._cmaps[key].copy() + value = self._cmaps[key] + + # Lazy loading from file + if isinstance(value, dict) and value.get("is_lazy"): + value = self._load_and_register_cmap(key, value) + if not value: + raise KeyError(f"Failed to load colormap {key!r} from file.") + + # Lazy loading for builtin matplotlib cmaps + if not isinstance(value, (ContinuousColormap, DiscreteColormap)): + value = _translate_cmap(value) + self._cmaps[key] = value + + value = value.copy() else: raise KeyError( f"Invalid colormap or color cycle name {key!r}. Options are: " @@ -3204,6 +3248,7 @@ def __getitem__(self, key): value = value.shifted(180) return value + @override def register(self, cmap, *, name=None, force=False): """ Add the colormap after validating and converting. @@ -3220,7 +3265,19 @@ def register(self, cmap, *, name=None, force=False): # surpress warning if the colormap is not generate by ultraplot if name not in self._builtin_cmaps: print(f"Overwriting {name!r} that was already registered") - self._cmaps[name] = cmap.copy() + self._cmaps[name] = cmap.copy(name=name) + + def register_lazy(self, name, path, type, is_default=False): + """ + Register a colormap to be loaded lazily from a file. + """ + name = self._translate_key(name, mirror=False) + self._cmaps[name] = { + "path": path, + "type": type, + "is_default": is_default, + "is_lazy": True, + } # Initialize databases diff --git a/ultraplot/config.py b/ultraplot/config.py index dc0458d6d..1b98da4ad 100644 --- a/ultraplot/config.py +++ b/ultraplot/config.py @@ -497,12 +497,11 @@ def register_cmaps(*args, user=None, local=None, default=False): for i, path in _iter_data_objects( "cmaps", *paths, user=user, local=local, default=default ): - cmap = pcolors.ContinuousColormap.from_file(path, warn_on_failure=True) - if not cmap: - continue - if i == 0 and cmap.name.lower() in pcolors.CMAPS_CYCLIC: - cmap.set_cyclic(True) - pcolors._cmap_database.register(cmap, name=cmap.name) + name, ext = os.path.splitext(os.path.basename(path)) + if ext and ext[1:] in ("json", "txt", "rgb", "xml", "hex"): + pcolors._cmap_database.register_lazy( + name, path, "continuous", is_default=(i == 0) + ) @docstring._snippet_manager @@ -541,10 +540,9 @@ def register_cycles(*args, user=None, local=None, default=False): for _, path in _iter_data_objects( "cycles", *paths, user=user, local=local, default=default ): - cmap = pcolors.DiscreteColormap.from_file(path, warn_on_failure=True) - if not cmap: - continue - pcolors._cmap_database.register(cmap, name=cmap.name) + name, ext = os.path.splitext(os.path.basename(path)) + if ext and ext[1:] in ("hex", "rgb", "txt"): + pcolors._cmap_database.register_lazy(name, path, "discrete") @docstring._snippet_manager diff --git a/ultraplot/tests/test_colors.py b/ultraplot/tests/test_colors.py new file mode 100644 index 000000000..d68a8ebba --- /dev/null +++ b/ultraplot/tests/test_colors.py @@ -0,0 +1,129 @@ +import os +import pytest +import numpy as np +import matplotlib.colors as mcolors + +from ultraplot import colors as pcolors +from ultraplot import config + + +@pytest.fixture(autouse=True) +def setup_teardown(): + """ + Reset the colormap database before and after each test. + """ + # This ensures a clean state for each test. + # The singleton instance is replaced with a new one. + pcolors._cmap_database = pcolors._init_cmap_database() + config.register_cmaps(default=True) + config.register_cycles(default=True) + yield + + +def test_lazy_loading_builtin(): + """ + Test that built-in colormaps are lazy-loaded. + """ + # Before access, it should be a matplotlib colormap + cmap_raw = pcolors._cmap_database._cmaps["viridis"] + assert isinstance( + cmap_raw, + ( + pcolors.ContinuousColormap, + pcolors.DiscreteColormap, + mcolors.ListedColormap, + ), + ) + + # After access, it should be an ultraplot colormap + cmap_get = pcolors._cmap_database.get_cmap("viridis") + assert isinstance(cmap_get, pcolors.ContinuousColormap) + + # The internal representation should also be updated + cmap_raw_after = pcolors._cmap_database._cmaps["viridis"] + assert isinstance(cmap_raw_after, pcolors.ContinuousColormap) + + +def test_case_insensitivity(): + """ + Test that colormap lookup is case-insensitive. + """ + cmap1 = pcolors._cmap_database.get_cmap("ViRiDiS") + cmap2 = pcolors._cmap_database.get_cmap("viridis") + assert cmap1.name.lower().startswith("_viridis") + assert cmap2.name.lower().startswith("_viridis") + + +def test_reversed_shifted(): + """ + Test reversed and shifted colormaps. + """ + # Create a simple colormap to test the reversal logic + # This avoids dependency on the exact definition of 'viridis' in matplotlib + colors_list = [(1, 0, 0), (0, 1, 0), (0, 0, 1)] # Red, Green, Blue + test_cmap = pcolors.ContinuousColormap.from_list("test_cmap", colors_list) + pcolors._cmap_database.register(test_cmap) + + cmap = pcolors._cmap_database.get_cmap("test_cmap") + cmap_r = pcolors._cmap_database.get_cmap("test_cmap_r") + + # Check name + assert cmap_r.name == "_test_cmap_copy_r" + # Check colors + # Start of original should be end of reversed + assert np.allclose(cmap(0.0), cmap_r(1.0)) + # End of original should be start of reversed + assert np.allclose(cmap(1.0), cmap_r(0.0)) + # Middle should be the same + assert np.allclose(cmap(0.5)[:3], cmap_r(0.5)[:3][::-1]) + + +def test_grays_translation(): + """ + Test that 'Grays' is translated to 'greys'. + """ + cmap_grays = pcolors._cmap_database.get_cmap("Grays") + assert cmap_grays.name.lower().startswith("_greys") + + +def test_lazy_loading_file(tmp_path): + """ + Test that colormaps from files are lazy-loaded. + """ + # Create a dummy colormap file + cmap_data = "1, 0, 0\n0, 1, 0\n0, 0, 1" + cmap_file = tmp_path / "my_test_cmap.rgb" + cmap_file.write_text(cmap_data) + + # Register it lazily + pcolors._cmap_database.register_lazy("my_test_cmap", str(cmap_file), "continuous") + + # Before access, it should be a lazy-load dict + cmap_raw = pcolors._cmap_database._cmaps["my_test_cmap"] + assert isinstance(cmap_raw, dict) + assert cmap_raw["is_lazy"] + + # After access, it should be an ultraplot colormap + cmap_get = pcolors._cmap_database.get_cmap("my_test_cmap") + assert isinstance(cmap_get, pcolors.ContinuousColormap) + assert cmap_get.name.lower().startswith("_my_test_cmap") + + # The internal representation should also be updated + cmap_raw_after = pcolors._cmap_database._cmaps["my_test_cmap"] + assert isinstance(cmap_raw_after, pcolors.ContinuousColormap) + + +def test_register_new(): + """ + Test registering a new colormap. + """ + colors_list = [(0, 0, 0), (1, 1, 1)] + new_cmap = pcolors.DiscreteColormap(colors_list, name="my_new_cmap") + pcolors._cmap_database.register(new_cmap) + + # Check it was registered + cmap_get = pcolors._cmap_database.get_cmap("my_new_cmap") + assert cmap_get.name.lower().startswith( + "_my_new_cmap" + ), f"Received {cmap_get.name.lower()} expected _my_new_cmap" + assert len(cmap_get.colors) == 2