Merge pull request #20227 from greglucas/cm-equals

tacaswell · web-flow · commit 136dc16a6c38 · 2021-05-14T18:52:24.000-04:00
ENH: Adding an equals method to colormaps
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -690,6 +690,17 @@ def __copy__(self):
         cmapobject._global = False
         return cmapobject
 
+    def __eq__(self, other):
+        if (not isinstance(other, Colormap) or self.name != other.name or
+                self.colorbar_extend != other.colorbar_extend):
+            return False
+        # To compare lookup tables the Colormaps have to be initialized
+        if not self._isinit:
+            self._init()
+        if not other._isinit:
+            other._init()
+        return np.array_equal(self._lut, other._lut)
+
     def get_bad(self):
         """Get the color for masked values."""
         if not self._isinit:
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -162,6 +162,29 @@ def test_colormap_copy():
     assert_array_equal(ret1, ret2)
 
 
+def test_colormap_equals():
+    cmap = plt.get_cmap("plasma")
+    cm_copy = cmap.copy()
+    # different object id's
+    assert cm_copy is not cmap
+    # But the same data should be equal
+    assert cm_copy == cmap
+    # Change the copy
+    cm_copy.set_bad('y')
+    assert cm_copy != cmap
+    # Make sure we can compare different sizes without failure
+    cm_copy._lut = cm_copy._lut[:10, :]
+    assert cm_copy != cmap
+    # Test different names are not equal
+    cm_copy = cmap.copy()
+    cm_copy.name = "Test"
+    assert cm_copy != cmap
+    # Test colorbar extends
+    cm_copy = cmap.copy()
+    cm_copy.colorbar_extend = not cmap.colorbar_extend
+    assert cm_copy != cmap
+
+
 def test_colormap_endian():
     """
     GitHub issue #1005: a bug in putmask caused erroneous
