ENH: add a copy method to colormaps

jklymak · jklymak · commit 9512d323fafb · 2021-03-07T13:11:32.000-08:00
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -561,7 +561,7 @@ def _warn_if_global_cmap_modified(cmap):
                     "colormap. In future versions, you will not be able to "
                     "modify a registered colormap in-place. To remove this "
                     "warning, you can make a copy of the colormap first. "
-                    f'cmap = copy.copy(mpl.cm.get_cmap("{cmap.name}"))'
+                    f'cmap = mpl.cm.get_cmap("{cmap.name}").copy()'
         )
 
 
@@ -845,6 +845,12 @@ def color_block(color):
                 f'over {color_block(self.get_over())}'
                 '</div>')
 
+    def copy(self):
+        """
+        Return a copy of the colormap.
+        """
+        return self.__copy__()
+
 
 class LinearSegmentedColormap(Colormap):
     """
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -150,6 +150,16 @@ def test_colormap_copy():
     with np.errstate(invalid='ignore'):
         ret2 = copied_cmap([-1, 0, .5, 1, np.nan, np.inf])
     assert_array_equal(ret1, ret2)
+    # again with the .copy method:
+    cmap = plt.cm.Reds
+    copied_cmap = cmap.copy()
+    with np.errstate(invalid='ignore'):
+        ret1 = copied_cmap([-1, 0, .5, 1, np.nan, np.inf])
+    cmap2 = copy.copy(copied_cmap)
+    cmap2.set_bad('g')
+    with np.errstate(invalid='ignore'):
+        ret2 = copied_cmap([-1, 0, .5, 1, np.nan, np.inf])
+    assert_array_equal(ret1, ret2)
 
 
 def test_colormap_endian():
