FIX: add missing method to ColormapRegistry

tacaswell · tacaswell · commit 329983d474d5 · 2022-10-06T18:08:27.000-04:00
After putting pending deprecations on `cm.get_cmap` we discovered that downstream libraries (pandas) were using the deprecated method to normalize between `None` (to get the default colormap), strings, and Colormap instances. This adds a method to `ColormapRegistry` to do this normalization. This can not replace our internal helper due to variations in what exceptions are raised. Closes #23981
diff --git a/doc/api/prev_api_changes/api_changes_3.6.0/deprecations.rst b/doc/api/prev_api_changes/api_changes_3.6.0/deprecations.rst
@@ -52,7 +52,13 @@ In Matplotlib 3.6 we have marked those top level functions as pending
 deprecation with the intention of deprecation in Matplotlib 3.7. The following
 functions have been marked for pending deprecation:
 
-- ``matplotlib.cm.get_cmap``; use ``matplotlib.colormaps[name]`` instead
+- ``matplotlib.cm.get_cmap``; use ``matplotlib.colormaps[name]`` instead if you
+  have a `str`.
+
+  **Added 3.6.1** Use `matplotlib.colors.ensure_colormap` if you have a string,
+  `None` or a `matplotlib.colors.Colormap` object that you want to convert to a
+  `matplotlib.colors.Colormap` instance.  Raises `KeyError` rather than
+  `ValueError` for missing strings.
 - ``matplotlib.cm.register_cmap``; use `matplotlib.colormaps.register
   <.ColormapRegistry.register>` instead
 - ``matplotlib.cm.unregister_cmap``; use `matplotlib.colormaps.unregister
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -193,6 +193,49 @@ def unregister(self, name):
                              "colormap.")
         self._cmaps.pop(name, None)
 
+    def ensure_colormap(self, cmap, *, default=None):
+        """
+        Ensure that at given object is a converted to a color map.
+
+        If *cmap* in `None`, returns *default*
+
+        Parameters
+        ----------
+        cmap : str, Colormap, None
+
+            - if a `Colormap`, return it
+            - if a string, look it up in mpl.colormaps
+            - if None, look up the default color map in mpl.colormaps
+
+        default : str, Colormap, None, optional
+            The default color map be returned if *cmap* in `None`
+
+            If `None` defaults to :rc:`image.cmap`
+
+        Returns
+        -------
+        Colormap
+
+        Raises
+        ------
+        KeyError
+        """
+        if cmap is None:
+            # if we have to fallback, call ourselves with the default value
+            return self.ensure_colormap(
+                # if the default default, use rcparams['image.cmap']
+                default
+                if default is not None
+                else mpl.rcParams["image.cmap"]
+            )
+
+        # if the user passed in a Colormap, simply return it
+        if isinstance(cmap, colors.Colormap):
+            return cmap
+
+        # otherwise, it must be a string so look it up
+        return self[cmap]
+
 
 # public access to the colormaps should be via `matplotlib.colormaps`. For now,
 # we still create the registry here, but that should stay an implementation
@@ -281,7 +324,12 @@ def _get_cmap(name=None, lut=None):
 # pyplot.
 get_cmap = _api.deprecated(
     '3.6',
-    name='get_cmap', pending=True, alternative="``matplotlib.colormaps[name]``"
+    name='get_cmap',
+    pending=True,
+    alternative=(
+        "``matplotlib.colormaps[name]`` " +
+        "or ``matplotlib.colormaps.ensure_colormap(obj)``"
+    )
 )(_get_cmap)
 
 
@@ -687,6 +735,8 @@ def _ensure_cmap(cmap):
     """
     Ensure that we have a `.Colormap` object.
 
+    For internal use to preserve type stability of errors.
+
     Parameters
     ----------
     cmap : None, str, Colormap
@@ -698,6 +748,7 @@ def _ensure_cmap(cmap):
     Returns
     -------
     Colormap
+
     """
     if isinstance(cmap, colors.Colormap):
         return cmap
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -109,6 +109,24 @@ def test_register_cmap():
             cm.register_cmap('nome', cmap='not a cmap')
 
 
+def test_ensure_cmap():
+    cr = mpl.colormaps
+    new_cm = mcolors.ListedColormap(cr["viridis"].colors, name='v2')
+
+    # check None, str, and Colormap pass
+    assert cr.ensure_colormap('plasma') == cr["plasma"]
+    assert cr.ensure_colormap(cr["magma"]) == cr["magma"]
+
+    # check default default
+    assert cr.ensure_colormap(None) == cr[mpl.rcParams['image.cmap']]
+    # check setting default
+    assert cr.ensure_colormap(None, default=new_cm) == new_cm
+    assert cr.ensure_colormap(None, default='magma') == cr["magma"]
+    bad_cmap = 'AardvarksAreAwkward'
+    with pytest.raises(KeyError, match=bad_cmap):
+        cr.ensure_colormap(bad_cmap)
+
+
 def test_double_register_builtin_cmap():
     name = "viridis"
     match = f"Re-registering the builtin cmap {name!r}."
