FIX: Update the callback in scalar mappable on norm change

greglucas · greglucas · commit 2dde636496b2 · 2021-05-24T20:14:52.000-06:00
When changing a ScalarMappable's norm, the previous norm
callback wasn't being removed and updated with the new one. This
adds that capability to the norm setter.
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -257,16 +257,13 @@ def __init__(self, norm=None, cmap=None):
             The colormap used to map normalized data values to RGBA colors.
         """
         self._A = None
-        self.norm = None  # So that the setter knows we're initializing.
+        self._norm = None  # So that the setter knows we're initializing.
         self.set_norm(norm)  # The Normalize instance of this ScalarMappable.
         self.cmap = None  # So that the setter knows we're initializing.
         self.set_cmap(cmap)  # The Colormap instance of this ScalarMappable.
         #: The last colorbar associated with this ScalarMappable. May be None.
         self.colorbar = None
         self.callbacksSM = cbook.CallbackRegistry()
-        # Connect to the Norm's callback
-        self._id_norm = self.norm.callbacksNorm.connect('changed',
-                                                        self.changed)
 
     def _scale_norm(self, norm, vmin, vmax):
         """
@@ -454,6 +451,30 @@ def set_cmap(self, cmap):
         if not in_init:
             self.changed()  # Things are not set up properly yet.
 
+    @property
+    def norm(self):
+        return self._norm
+
+    @norm.setter
+    def norm(self, norm):
+        _api.check_isinstance((colors.Normalize, None), norm=norm)
+        if norm is None:
+            norm = colors.Normalize()
+
+        if norm is self.norm:
+            # We aren't updating anything
+            return
+
+        in_init = self.norm is None
+        # Remove the current callback and connect to the new one
+        if not in_init:
+            self.norm.callbacksNorm.disconnect(self._id_norm)
+        self._norm = norm
+        self._id_norm = self.norm.callbacksNorm.connect('changed',
+                                                        self.changed)
+        if not in_init:
+            self.changed()
+
     def set_norm(self, norm):
         """
         Set the normalization instance.
@@ -468,13 +489,7 @@ def set_norm(self, norm):
         the norm of the mappable will reset the norm, locator, and formatters
         on the colorbar to default.
         """
-        _api.check_isinstance((colors.Normalize, None), norm=norm)
-        in_init = self.norm is None
-        if norm is None:
-            norm = colors.Normalize()
         self.norm = norm
-        if not in_init:
-            self.changed()  # Things are not set up properly yet.
 
     def autoscale(self):
         """
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -1450,13 +1450,15 @@ def test_scalarmappable_norm_update():
     norm.clip = True
     assert sm.stale
     # change to the CenteredNorm and TwoSlopeNorm to test those
+    # Also make sure that updating the norm directly and with
+    # set_norm both update the Norm callback
     norm = mcolors.CenteredNorm()
-    sm = matplotlib.cm.ScalarMappable(norm=norm, cmap='plasma')
+    sm.norm = norm
     sm.stale = False
     norm.vcenter = 1
     assert sm.stale
     norm = mcolors.TwoSlopeNorm(vcenter=0, vmin=-1, vmax=1)
-    sm = matplotlib.cm.ScalarMappable(norm=norm, cmap='plasma')
+    sm.set_norm(norm)
     sm.stale = False
     norm.vcenter = 1
     assert sm.stale
