Merge pull request #8375 from vidursatija/i_8299

tacaswell · web-flow · commit d50843e70378 · 2017-05-24T22:49:46.000-04:00
FIX: implemented copy on color map objects
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -530,6 +530,16 @@ def __call__(self, X, alpha=None, bytes=False):
             rgba = tuple(rgba[0, :])
         return rgba
 
+    def __copy__(self):
+        """Create new object with the same class, update attributes
+        """
+        cls = self.__class__
+        cmapobject = cls.__new__(cls)
+        cmapobject.__dict__.update(self.__dict__)
+        if self._isinit:
+            cmapobject._lut = np.copy(self._lut)
+        return cmapobject
+
     def set_bad(self, color='k', alpha=None):
         """Set color to be used for masked values.
         """
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -1,6 +1,7 @@
 from __future__ import (absolute_import, division, print_function,
                         unicode_literals)
 
+import copy
 import six
 import itertools
 import warnings
@@ -43,6 +44,18 @@ def test_resample():
     assert_array_almost_equal(lc3([0, 0.5, 1]), expected)
 
 
+def test_colormap_copy():
+    cm = plt.cm.Reds
+    cm_copy = copy.copy(cm)
+    with np.errstate(invalid='ignore'):
+        ret1 = cm_copy([-1, 0, .5, 1, np.nan, np.inf])
+    cm2 = copy.copy(cm_copy)
+    cm2.set_bad('g')
+    with np.errstate(invalid='ignore'):
+        ret2 = cm_copy([-1, 0, .5, 1, np.nan, np.inf])
+    assert_array_equal(ret1, ret2)
+
+
 def test_colormap_endian():
     """
     Github issue #1005: a bug in putmask caused erroneous
