MAINT: Adding an equals method to colormaps

greglucas · greglucas · commit f94d64d0ff04 · 2021-05-13T21:53:56.000-06:00
This adds an __eq__ method to Colormap to test the lookup
tables for equality.
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -690,6 +690,16 @@ def __copy__(self):
         cmapobject._global = False
         return cmapobject
 
+    def __eq__(self, other):
+        if not isinstance(other, Colormap):
+            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,21 @@ 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
+
+
 def test_colormap_endian():
     """
     GitHub issue #1005: a bug in putmask caused erroneous
