Merge pull request #15127 from tacaswell/register_cmap_test

dopplershift · web-flow · commit 6aca58edbd88 · 2020-09-25T12:36:54.000-06:00
ENH/API: improvements to register_cmap
diff --git a/doc/api/next_api_changes/behavior/15127-TAC.rst b/doc/api/next_api_changes/behavior/15127-TAC.rst
@@ -0,0 +1,8 @@
+Raise or warn on registering a colormap twice
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When using `matplotlib.cm.register_cmap` to register a user provided
+or third-party colormap it will now raise a `ValueError` if trying to
+over-write one of the built in colormaps and warn if trying to over
+write a user registered colormap.  This may raise for user-registered
+colormaps in the future.
diff --git a/doc/users/next_whats_new/2019-08_tac.rst b/doc/users/next_whats_new/2019-08_tac.rst
@@ -0,0 +1,6 @@
+
+Add ``cm.unregister_cmap`` function
+-----------------------------------
+
+`.cm.unregister_cmap` allows users to remove a colormap that they
+have previously registered.
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -24,6 +24,7 @@
 from matplotlib import _api, colors, cbook
 from matplotlib._cm import datad
 from matplotlib._cm_listed import cmaps as cmaps_listed
+from matplotlib.cbook import _warn_external
 
 
 LUTSIZE = mpl.rcParams['image.lut']
@@ -95,37 +96,46 @@ def _warn_deprecated(self):
 locals().update(_cmap_registry)
 # This is no longer considered public API
 cmap_d = _DeprecatedCmapDictWrapper(_cmap_registry)
-
+__builtin_cmaps = tuple(_cmap_registry)
 
 # Continue with definitions ...
 
 
-def register_cmap(name=None, cmap=None, data=None, lut=None):
+def register_cmap(name=None, cmap=None, *, override_builtin=False):
     """
     Add a colormap to the set recognized by :func:`get_cmap`.
 
-    It can be used in two ways::
+    Register a new colormap to be accessed by name ::
+
+        LinearSegmentedColormap('swirly', data, lut)
+        register_cmap(cmap=swirly_cmap)
 
-        register_cmap(name='swirly', cmap=swirly_cmap)
+    Parameters
+    ----------
+    name : str, optional
+       The name that can be used in :func:`get_cmap` or :rc:`image.cmap`
 
-        register_cmap(name='choppy', data=choppydata, lut=128)
+       If absent, the name will be the :attr:`~matplotlib.colors.Colormap.name`
+       attribute of the *cmap*.
 
-    In the first case, *cmap* must be a :class:`matplotlib.colors.Colormap`
-    instance.  The *name* is optional; if absent, the name will
-    be the :attr:`~matplotlib.colors.Colormap.name` attribute of the *cmap*.
+    cmap : matplotlib.colors.Colormap
+       Despite being the second argument and having a default value, this
+       is a required argument.
 
-    The second case is deprecated. Here, the three arguments are passed to
-    the :class:`~matplotlib.colors.LinearSegmentedColormap` initializer,
-    and the resulting colormap is registered. Instead of this implicit
-    colormap creation, create a `.LinearSegmentedColormap` and use the first
-    case: ``register_cmap(cmap=LinearSegmentedColormap(name, data, lut))``.
+    override_builtin : bool
+
+        Allow built-in colormaps to be overridden by a user-supplied
+        colormap.
+
+        Please do not use this unless you are sure you need it.
 
     Notes
     -----
     Registering a colormap stores a reference to the colormap object
     which can currently be modified and inadvertantly change the global
     colormap state. This behavior is deprecated and in Matplotlib 3.5
     the registered colormap will be immutable.
+
     """
     cbook._check_isinstance((str, None), name=name)
     if name is None:
@@ -134,23 +144,21 @@ def register_cmap(name=None, cmap=None, data=None, lut=None):
         except AttributeError as err:
             raise ValueError("Arguments must include a name or a "
                              "Colormap") from err
-    if isinstance(cmap, colors.Colormap):
-        cmap._global = True
-        _cmap_registry[name] = cmap
-        return
-    if lut is not None or data is not None:
-        cbook.warn_deprecated(
-            "3.3",
-            message="Passing raw data via parameters data and lut to "
-                    "register_cmap() is deprecated since %(since)s and will "
-                    "become an error %(removal)s. Instead use: register_cmap("
-                    "cmap=LinearSegmentedColormap(name, data, lut))")
-    # For the remainder, let exceptions propagate.
-    if lut is None:
-        lut = mpl.rcParams['image.lut']
-    cmap = colors.LinearSegmentedColormap(name, data, lut)
+    if name in _cmap_registry:
+        if not override_builtin and name in __builtin_cmaps:
+            msg = f"Trying to re-register the builtin cmap {name!r}."
+            raise ValueError(msg)
+        else:
+            msg = f"Trying to register the cmap {name!r} which already exists."
+            _warn_external(msg)
+
+    if not isinstance(cmap, colors.Colormap):
+        raise ValueError("You must pass a Colormap instance. "
+                         f"You passed {cmap} a {type(cmap)} object.")
+
     cmap._global = True
     _cmap_registry[name] = cmap
+    return
 
 
 def get_cmap(name=None, lut=None):
@@ -187,6 +195,47 @@ def get_cmap(name=None, lut=None):
         return _cmap_registry[name]._resample(lut)
 
 
+def unregister_cmap(name):
+    """
+    Remove a colormap recognized by :func:`get_cmap`.
+
+    You may not remove built-in colormaps.
+
+    If the named colormap is not registered, returns with no error, raises
+    if you try to de-register a default colormap.
+
+    .. warning ::
+
+      Colormap names are currently a shared namespace that may be used
+      by multiple packages. Use `unregister_cmap` only if you know you
+      have registered that name before. In particular, do not
+      unregister just in case to clean the name before registering a
+      new colormap.
+
+    Parameters
+    ----------
+    name : str
+        The name of the colormap to be un-registered
+
+    Returns
+    -------
+    ColorMap or None
+        If the colormap was registered, return it if not return `None`
+
+    Raises
+    ------
+    ValueError
+       If you try to de-register a default built-in colormap.
+
+    """
+    if name not in _cmap_registry:
+        return
+    if name in __builtin_cmaps:
+        raise ValueError(f"cannot unregister {name!r} which is a builtin "
+                         "colormap.")
+    return _cmap_registry.pop(name)
+
+
 class ScalarMappable:
     """
     A mixin class to map scalar data to RGBA.
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -652,7 +652,7 @@ def get_bad(self):
         """Get the color for masked values."""
         if not self._isinit:
             self._init()
-        return self._lut[self._i_bad]
+        return np.array(self._lut[self._i_bad])
 
     def set_bad(self, color='k', alpha=None):
         """Set the color for masked values."""
@@ -665,7 +665,7 @@ def get_under(self):
         """Get the color for low out-of-range values."""
         if not self._isinit:
             self._init()
-        return self._lut[self._i_under]
+        return np.array(self._lut[self._i_under])
 
     def set_under(self, color='k', alpha=None):
         """Set the color for low out-of-range values."""
@@ -678,7 +678,7 @@ def get_over(self):
         """Get the color for high out-of-range values."""
         if not self._isinit:
             self._init()
-        return self._lut[self._i_over]
+        return np.array(self._lut[self._i_over])
 
     def set_over(self, color='k', alpha=None):
         """Set the color for high out-of-range values."""
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -64,14 +64,44 @@ def test_resample():
 
 
 def test_register_cmap():
-    new_cm = copy.copy(plt.cm.viridis)
-    cm.register_cmap('viridis2', new_cm)
-    assert plt.get_cmap('viridis2') == new_cm
+    new_cm = copy.copy(cm.get_cmap("viridis"))
+    target = "viridis2"
+    cm.register_cmap(target, new_cm)
+    assert plt.get_cmap(target) == new_cm
 
     with pytest.raises(ValueError,
-                       match='Arguments must include a name or a Colormap'):
+                       match="Arguments must include a name or a Colormap"):
         cm.register_cmap()
 
+    with pytest.warns(UserWarning):
+        cm.register_cmap(target, new_cm)
+
+    cm.unregister_cmap(target)
+    with pytest.raises(ValueError,
+                       match=f'{target!r} is not a valid value for name;'):
+        cm.get_cmap(target)
+    # test that second time is error free
+    cm.unregister_cmap(target)
+
+    with pytest.raises(ValueError, match="You must pass a Colormap instance."):
+        cm.register_cmap('nome', cmap='not a cmap')
+
+
+def test_double_register_builtin_cmap():
+    name = "viridis"
+    match = f"Trying to re-register the builtin cmap {name!r}."
+    with pytest.raises(ValueError, match=match):
+        cm.register_cmap(name, cm.get_cmap(name))
+    with pytest.warns(UserWarning):
+        cm.register_cmap(name, cm.get_cmap(name), override_builtin=True)
+
+
+def test_unregister_builtin_cmap():
+    name = "viridis"
+    match = f'cannot unregister {name!r} which is a builtin colormap.'
+    with pytest.raises(ValueError, match=match):
+        cm.unregister_cmap(name)
+
 
 def test_colormap_global_set_warn():
     new_cm = plt.get_cmap('viridis')
@@ -94,7 +124,8 @@ def test_colormap_global_set_warn():
         new_cm.set_under('k')
 
     # Re-register the original
-    plt.register_cmap(cmap=orig_cmap)
+    with pytest.warns(UserWarning):
+        plt.register_cmap(cmap=orig_cmap, override_builtin=True)
 
 
 def test_colormap_dict_deprecate():
@@ -1187,6 +1218,16 @@ def test_get_under_over_bad():
     assert_array_equal(cmap.get_bad(), cmap(np.nan))
 
 
+@pytest.mark.parametrize('kind', ('over', 'under', 'bad'))
+def test_non_mutable_get_values(kind):
+    cmap = copy.copy(plt.get_cmap('viridis'))
+    init_value = getattr(cmap, f'get_{kind}')()
+    getattr(cmap, f'set_{kind}')('k')
+    black_value = getattr(cmap, f'get_{kind}')()
+    assert np.all(black_value == [0, 0, 0, 1])
+    assert not np.all(init_value == black_value)
+
+
 def test_colormap_alpha_array():
     cmap = plt.get_cmap('viridis')
     vals = [-1, 0.5, 2]  # under, valid, over
