Move all Colormap extremes setter logic into a single _set_extremes()

Author: timhoffm

lib/matplotlib/colors.py

@@ -877,9 +877,7 @@ def get_bad(self):
 
     def set_bad(self, color='k', alpha=None):
         """Set the color for masked values."""
-        self._rgba_bad = to_rgba(color, alpha)
-        if self._isinit:
-            self._set_extremes()
+        self._set_extremes(bad=(color, alpha))
 
     def get_under(self):
         """Get the color for low out-of-range values."""
@@ -889,9 +887,7 @@ def get_under(self):
 
     def set_under(self, color='k', alpha=None):
         """Set the color for low out-of-range values."""
-        self._rgba_under = to_rgba(color, alpha)
-        if self._isinit:
-            self._set_extremes()
+        self._set_extremes(under=(color, alpha))
 
     def get_over(self):
         """Get the color for high out-of-range values."""
@@ -901,21 +897,14 @@ def get_over(self):
 
     def set_over(self, color='k', alpha=None):
         """Set the color for high out-of-range values."""
-        self._rgba_over = to_rgba(color, alpha)
-        if self._isinit:
-            self._set_extremes()
+        self._set_extremes(over=(color, alpha))
 
     def set_extremes(self, *, bad=None, under=None, over=None):
         """
         Set the colors for masked (*bad*) values and, when ``norm.clip =
         False``, low (*under*) and high (*over*) out-of-range values.
         """
-        if bad is not None:
-            self.set_bad(bad)
-        if under is not None:
-            self.set_under(under)
-        if over is not None:
-            self.set_over(over)
+        self._set_extremes(bad=bad, under=under, over=over)
 
     def with_extremes(self, *, bad=None, under=None, over=None):
         """
@@ -924,19 +913,33 @@ def with_extremes(self, *, bad=None, under=None, over=None):
         out-of-range values, have been set accordingly.
         """
         new_cm = self.copy()
-        new_cm.set_extremes(bad=bad, under=under, over=over)
+        new_cm._set_extremes(bad=bad, under=under, over=over)
         return new_cm
 
-    def _set_extremes(self):
-        if self._rgba_under:
-            self._lut[self._i_under] = self._rgba_under
-        else:
-            self._lut[self._i_under] = self._lut[0]
-        if self._rgba_over:
-            self._lut[self._i_over] = self._rgba_over
-        else:
-            self._lut[self._i_over] = self._lut[self.N - 1]
-        self._lut[self._i_bad] = self._rgba_bad
+    def _set_extremes(self, bad=None, under=None, over=None):
+        """
+        Set the colors for masked (*bad*) and out-of-range (*under* and *over) values.
+
+        Parameters that are None are left unchanged.
+        """
+        if bad is not None:
+            self._rgba_bad = to_rgba(bad)
+        if under is not None:
+            self._rgba_under = to_rgba(under)
+        if over is not None:
+            self._rgba_over = to_rgba(over)
+
+        # ensure that an existing _lut is updated
+        if self._isinit:
+            if self._rgba_under:
+                self._lut[self._i_under] = self._rgba_under
+            else:
+                self._lut[self._i_under] = self._lut[0]
+            if self._rgba_over:
+                self._lut[self._i_over] = self._rgba_over
+            else:
+                self._lut[self._i_over] = self._lut[self.N - 1]
+            self._lut[self._i_bad] = self._rgba_bad
 
     def with_alpha(self, alpha):
         """
@@ -1154,7 +1157,7 @@ def _init(self):
             self._lut[:-3, 3] = _create_lookup_table(
                 self.N, self._segmentdata['alpha'], 1)
         self._isinit = True
-        self._set_extremes()
+        self._update_extremes()
 
     def set_gamma(self, gamma):
         """Set a new gamma value and regenerate colormap."""
@@ -1346,7 +1349,7 @@ def _init(self):
         self._lut = np.zeros((self.N + 3, 4), float)
         self._lut[:-3] = to_rgba_array(self.colors)
         self._isinit = True
-        self._set_extremes()
+        self._update_extremes()
 
     @property
     def monochrome(self):