ENH/API: be more defensive in cm.register_cmap + deregister

tacaswell · dstansby · tacaswell · commit 1a9dacff1cb0 · 2020-09-25T09:22:42.000-04:00
- raise if you try to over write a default colormap
 - warn if you try to over write a user-added colormap
 - add method to de-register a color map
 - add escape hatch to force re-registering builtin colormaps

Co-Authored-By: David Stansby &lt;dstansby@gmail.com&gt;
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,12 +96,12 @@ 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):
+def register_cmap(name=None, cmap=None, *, override_builtin=False):
     """
     Add a colormap to the set recognized by :func:`get_cmap`.
 
@@ -121,13 +122,20 @@ def register_cmap(name=None, cmap=None):
        Despite being the second argument and having a default value, this
        is a required argument.
 
+    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:
@@ -136,10 +144,18 @@ def register_cmap(name=None, cmap=None):
         except AttributeError as err:
             raise ValueError("Arguments must include a name or a "
                              "Colormap") from err
+    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
@@ -179,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/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -64,18 +64,45 @@ 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')
     # Store the old value so we don't override the state later on.
@@ -97,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():
