Warning when modifying the state of a global colormap.

greglucas · greglucas · commit 67ffc664bc9b · 2020-03-31T20:50:25.000-06:00
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -65,12 +65,16 @@ def _gen_cmap_d():
     # Generate reversed cmaps.
     for cmap in list(cmap_d.values()):
         rmap = cmap.reversed()
+        cmap._global = True
+        rmap._global = True
         cmap_d[rmap.name] = rmap
     return cmap_d
 
 
-cmap_d = _gen_cmap_d()
-locals().update(cmap_d)
+_cmap_d = _gen_cmap_d()
+locals().update(_cmap_d)
+# This is no longer considered public API
+cmap_d = _cmap_d
 
 
 # Continue with definitions ...
@@ -104,7 +108,8 @@ def register_cmap(name=None, cmap=None, data=None, lut=None):
             raise ValueError("Arguments must include a name or a "
                              "Colormap") from err
     if isinstance(cmap, colors.Colormap):
-        cmap_d[name] = cmap
+        cmap._global = True
+        _cmap_d[name] = cmap
         return
     if lut is not None or data is not None:
         cbook.warn_deprecated(
@@ -117,7 +122,8 @@ def register_cmap(name=None, cmap=None, data=None, lut=None):
     if lut is None:
         lut = mpl.rcParams['image.lut']
     cmap = colors.LinearSegmentedColormap(name, data, lut)
-    cmap_d[name] = cmap
+    cmap._global = True
+    _cmap_d[name] = cmap
 
 
 def get_cmap(name=None, lut=None):
@@ -130,9 +136,12 @@ def get_cmap(name=None, lut=None):
     Parameters
     ----------
     name : `matplotlib.colors.Colormap` or str or None, default: None
-        If a `.Colormap` instance, it will be returned.  Otherwise, the name of
-        a colormap known to Matplotlib, which will be resampled by *lut*.  The
-        default, None, means :rc:`image.cmap`.
+        If a `.Colormap` instance, it will be returned.
+        Otherwise, the name of a colormap known to Matplotlib, which will
+        be resampled by *lut*. Currently, this returns the global colormap
+        object, which is deprecated. In Matplotlib 3.5, you will no longer be
+        able to modify the global colormaps in-place.
+        The default, None, means :rc:`image.cmap`.
     lut : int or None, default: None
         If *name* is not already a Colormap instance and *lut* is not None, the
         colormap will be resampled to have *lut* entries in the lookup table.
@@ -141,11 +150,11 @@ def get_cmap(name=None, lut=None):
         name = mpl.rcParams['image.cmap']
     if isinstance(name, colors.Colormap):
         return name
-    cbook._check_in_list(sorted(cmap_d), name=name)
+    cbook._check_in_list(sorted(_cmap_d), name=name)
     if lut is None:
-        return cmap_d[name]
+        return _cmap_d[name]
     else:
-        return cmap_d[name]._resample(lut)
+        return _cmap_d[name]._resample(lut)
 
 
 class ScalarMappable:
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -484,6 +484,21 @@ def makeMappingArray(N, data, gamma=1.0):
     return _create_lookup_table(N, data, gamma)
 
 
+def _deprecate_global_cmap(cmap):
+    if hasattr(cmap, '_global') and cmap._global:
+        cbook.warn_deprecated(
+            "3.3",
+            message="You are modifying the state of a globally registered "
+                    "colormap. In future versions, you will not be able "
+                    "to modify a globally registered colormap directly. "
+                    "To eliminate this warning until then, you can make "
+                    "a copy of the requested colormap before modifying it. ",
+            alternative="To modify a colormap without overwriting the "
+                        "global state, you can make a copy of the colormap "
+                        f"first. cmap = copy.copy(mpl.cm.get_cmap({cmap.name})"
+        )
+
+
 class Colormap:
     """
     Baseclass for all scalar to RGBA mappings.
@@ -599,13 +614,15 @@ def __copy__(self):
         cmapobject.__dict__.update(self.__dict__)
         if self._isinit:
             cmapobject._lut = np.copy(self._lut)
+        cmapobject._global = False
         return cmapobject
 
     def set_bad(self, color='k', alpha=None):
         """Set the color for masked values."""
         self._rgba_bad = to_rgba(color, alpha)
         if self._isinit:
             self._set_extremes()
+        _deprecate_global_cmap(self)
 
     def set_under(self, color='k', alpha=None):
         """
@@ -614,6 +631,7 @@ def set_under(self, color='k', alpha=None):
         self._rgba_under = to_rgba(color, alpha)
         if self._isinit:
             self._set_extremes()
+        _deprecate_global_cmap(self)
 
     def set_over(self, color='k', alpha=None):
         """
@@ -622,6 +640,7 @@ def set_over(self, color='k', alpha=None):
         self._rgba_over = to_rgba(color, alpha)
         if self._isinit:
             self._set_extremes()
+        _deprecate_global_cmap(self)
 
     def _set_extremes(self):
         if self._rgba_under:
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -70,6 +70,62 @@ def test_register_cmap():
         cm.register_cmap()
 
 
+@pytest.mark.skipif(matplotlib.__version__ < "3.5",
+                    reason="This test modifies the global state of colormaps.")
+def test_colormap_global_unchanged_state():
+    new_cm = plt.get_cmap('viridis')
+    new_cm.set_over('b')
+    # Make sure that this didn't mess with the original viridis cmap
+    assert new_cm != plt.get_cmap('viridis')
+
+    # Also test that pyplot access doesn't mess the original up
+    new_cm = plt.cm.viridis
+    new_cm.set_over('b')
+    assert new_cm != plt.get_cmap('viridis')
+
+
+@pytest.mark.skipif(matplotlib.__version__ < "3.4",
+                    reason="This test modifies the global state of colormaps.")
+def test_colormap_global_get_warn():
+    new_cm = plt.get_cmap('viridis')
+    # Store the old value so we don't override the state later on.
+    orig_cmap = copy.copy(new_cm)
+    new_cm.set_under('k')
+    with pytest.warns(cbook.MatplotlibDeprecationWarning,
+                      match="You are modifying the state of a globally"):
+        # This should warn now because we've modified the global state
+        # without registering it
+        plt.get_cmap('viridis')
+
+    # Test that re-registering the original cmap clears the warning
+    plt.register_cmap(cmap=orig_cmap)
+    plt.get_cmap('viridis')
+
+
+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.
+    orig_cmap = copy.copy(new_cm)
+    with pytest.warns(cbook.MatplotlibDeprecationWarning,
+                      match="You are modifying the state of a globally"):
+        # This should warn now because we've modified the global state
+        new_cm.set_under('k')
+
+    # This shouldn't warn because it is a copy
+    copy.copy(new_cm).set_under('b')
+
+    # Test that registering and then modifying warns
+    plt.register_cmap(name='test_cm', cmap=copy.copy(orig_cmap))
+    new_cm = plt.get_cmap('test_cm')
+    with pytest.warns(cbook.MatplotlibDeprecationWarning,
+                      match="You are modifying the state of a globally"):
+        # This should warn now because we've modified the global state
+        new_cm.set_under('k')
+
+    # Re-register the original
+    plt.register_cmap(cmap=orig_cmap)
+
+
 def test_colormap_copy():
     cm = plt.cm.Reds
     cm_copy = copy.copy(cm)
