Merge pull request #377 from neutrinoceros/cli/enh/support_colormap_e…

…xtensions_other_than_cblind

ENH: (CLI): support colormap extension packages other than cblind (cmasher, cmocean, cmyt)

nonos/main.py

@@ -11,12 +11,13 @@
 import re
 import time
 from collections import ChainMap
+from importlib import import_module
 from importlib.metadata import version
+from importlib.util import find_spec
 from multiprocessing import Pool
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
 
-import cblind  # noqa
 import inifix
 import numpy as np
 from packaging.version import Version
@@ -50,6 +51,13 @@
 NONOS_VERSION = version("nonos")
 INIFIX_GE_5_0 = Version(version("inifix")) >= Version("5.0.0")
 
+KNOWN_CMAP_PACKAGE_PREFIXES = {
+    "cb": "cblind",
+    "cmo": "cmocean",
+    "cmr": "cmasher",
+    "cmyt": "cmyt",
+}
+
 
 def get_non_interactive_figure(fmt: str) -> "Figure":
     from matplotlib.backends.backend_agg import FigureCanvasAgg
@@ -91,7 +99,7 @@ def process_field(
     vmin,
     vmax,
     scaling: float,
-    cmap,
+    cmap: str,
     title,
     unit_conversion: int,
     datadir,
@@ -155,32 +163,47 @@ def process_field(
     else:
         fig = get_non_interactive_figure(fmt)
 
+    plot_kwargs = {
+        "log": log,
+        "vmin": vmin,
+        "vmax": vmax,
+        "cmap": cmap,
+        "title": f"${title}$",
+        "unit_conversion": unit_conversion,
+    }
+
+    if cm_prefix := cmap.rpartition(".")[0]:
+        if (cm_package := KNOWN_CMAP_PACKAGE_PREFIXES.get(cm_prefix)) is None:
+            print_warn(
+                f"requested colormap {cmap!r} with the unknown prefix {cm_prefix!r}. "
+                "The default colormap will be used instead."
+            )
+            plot_kwargs.pop("cmap")
+        elif not find_spec(cm_package):
+            print_warn(
+                f"requested colormap {cmap!r}, but {cm_package} is not installed. "
+                "The default colormap will be used instead."
+            )
+            plot_kwargs.pop("cmap")
+        else:
+            # all known colormap packages work by registering their colormaps
+            # at import time.
+            import_module(cm_package)
+
     ax = fig.add_subplot(111, polar=False)
     if dim == 1:
-        dsop.map(plane[0], rotate_with=planet_file).plot(
-            fig,
-            ax,
-            log=log,
-            vmin=vmin,
-            vmax=vmax,
-            title=f"${title}$",
-            unit_conversion=unit_conversion,
-        )
+        if "cmap" in plot_kwargs:
+            plot_kwargs.pop("cmap")
+
+        dsop.map(plane[0], rotate_with=planet_file).plot(fig, ax, **plot_kwargs)
         akey = dsop.map(plane[0], rotate_with=planet_file).dict_plotable["abscissa"]
         avalue = dsop.map(plane[0], rotate_with=planet_file).dict_plotable[akey]
         extent = parse_range(extent, dim=dim)
         extent = range_converter(extent, abscissa=avalue, ordinate=np.zeros(2))
         ax.set_xlim(extent[0], extent[1])
     elif dim == 2:
         dsop.map(plane[0], plane[1], rotate_with=planet_file).plot(
-            fig,
-            ax,
-            log=log,
-            vmin=vmin,
-            vmax=vmax,
-            cmap=cmap,
-            title=f"${title}$",
-            unit_conversion=unit_conversion,
+            fig, ax, **plot_kwargs
         )
         akey = dsop.map(plane[0], plane[1], rotate_with=planet_file).dict_plotable[
             "abscissa"

requirements/tests_all.txt

@@ -1,2 +1,5 @@
 -r tests_min.txt
+cblind>=2.3.0
+cmocean>=3.0.3
+cmyt>=2.0.0
 numexpr>=2.8.3

tests/test_plotting.py

@@ -1,4 +1,5 @@
 import os
+from importlib.util import find_spec
 
 import numpy.testing as npt
 import pytest
@@ -199,3 +200,80 @@ def test_reg(test_data_dir, map_args):
     fig = Figure()
     ax = fig.add_subplot()
     ds["RHO"].map(*map_args).plot(fig, ax)
+
+
+@pytest.mark.parametrize(
+    "cmap",
+    [
+        pytest.param(
+            cmap_name,
+            marks=pytest.mark.skipif(
+                not find_spec(pkg_name), reason=f"{pkg_name} is not installed"
+            ),
+        )
+        for pkg_name, cmap_name in [
+            ("cblind", "cb.rainbow"),
+            ("cmocean", "cmo.thermal"),
+            ("cmasher", "cmr.dusk"),
+            ("cmyt", "cmyt.arbre"),
+        ]
+    ],
+)
+def test_colormap_extensions_integration(cmap, capsys, test_data_dir, tmp_path):
+    simdir = str(test_data_dir / "idefix_planet3d")
+    os.chdir(tmp_path)
+    ret = main(["-dir", simdir, "-geometry", "polar", "-cmap", cmap])
+    assert ret == 0
+    out, err = capsys.readouterr()
+    assert out == ""
+    assert err == ""
+
+
+@pytest.mark.parametrize(
+    "cmap, pkg",
+    [
+        pytest.param(
+            cmap_name,
+            pkg_name,
+            marks=pytest.mark.skipif(
+                find_spec(pkg_name), reason=f"{pkg_name} is installed"
+            ),
+        )
+        for pkg_name, cmap_name in [
+            ("cblind", "cb.rainbow"),
+            ("cmocean", "cmo.thermal"),
+            ("cmasher", "cmr.dusk"),
+            ("cmyt", "cmyt.arbre"),
+        ]
+    ],
+)
+def test_colormap_extensions_missing_package(
+    cmap, pkg, capsys, test_data_dir, tmp_path
+):
+    simdir = str(test_data_dir / "idefix_planet3d")
+    os.chdir(tmp_path)
+    ret = main(["-dir", simdir, "-geometry", "polar", "-cmap", cmap])
+    assert ret == 0
+    out, err = capsys.readouterr()
+    assert out == ""
+    assert err.replace("\n", "") == (
+        f"🦴 Warning requested colormap {cmap!r}, but {pkg} is not installed. "
+        "The default colormap will be used instead."
+    )
+
+
+def test_unknown_colormap_package_prefix(capsys, test_data_dir, tmp_path):
+    simdir = str(test_data_dir / "idefix_planet3d")
+    os.chdir(tmp_path)
+    ret = main(
+        ["-dir", simdir, "-geometry", "polar", "-cmap", "cmunknown.thismapdoesnexist"]
+    )
+    assert ret == 0
+
+    out, err = capsys.readouterr()
+    assert out == ""
+    assert err.replace("\n", "") == (
+        "🦴 Warning requested colormap 'cmunknown.thismapdoesnexist' "
+        "with the unknown prefix 'cmunknown'. "
+        "The default colormap will be used instead."
+    )