Proper behavior with extend kwarg

kdpenner · kdpenner · commit 3403cb03d715 · 2020-05-19T13:48:48.000-04:00
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -892,16 +892,38 @@ def _process_values(self, b=None):
             b = self.boundaries
         if b is not None:
             self._boundaries = np.asarray(b, dtype=float)
-            if self.values is None:
-                self._values = []
-                for i in range(len(self._boundaries)-1):
-                    self._values.extend(np.linspace(self._boundaries[i],
-                                        self._boundaries[i+1], 1000,
-                                        endpoint = False))
-                self._values.extend([self._boundaries[-1]])
-                self._values = np.array(self._values)
-                self._boundaries = self._values
+            if self.values is None and self.boundaries is not None:
+                # define the outside indices
+                if self.extend == 'min':
+                    i = 0
+                elif self.extend == 'max':
+                    i = -1
+                elif self.extend == 'both':
+                    i = [0, -1]
+                elif self.extend == 'neither':
+                    i = slice(None)
+                # add values to the colorbar insides
+                boundaries_inside = self._boundaries[self._inside]
+                self._values = np.linspace(boundaries_inside[:-1],
+                                           np.roll(boundaries_inside, -1)[:-1],
+                                           1000, endpoint = False,
+                                           axis = -1).flatten()
+                self._values = np.append(self._values, boundaries_inside[-1])
+                # add back in the boundaries for the extensions
+                to_add = np.atleast_1d(self._boundaries[i])
+                self._boundaries = np.sort(np.concatenate([to_add,
+                                                           self._values]))
+                # calculate values at which to plot colors
                 self._values = 0.5 * (self._values[:-1] + self._values[1:])
+                # add back in the boundaries for the extensions, display
+                # colors of extensions at boundary values
+                self._values = np.sort(np.concatenate([to_add,
+                                                       self._values]))
+                if isinstance(self.norm, colors.NoNorm):
+                    self._values = (self._values + 0.00001).astype(np.int16)
+            elif self.values is None and self.boundaries is None:
+                self._values = 0.5 * (self._boundaries[:-1]
+                                      + self._boundaries[1:])
                 if isinstance(self.norm, colors.NoNorm):
                     self._values = (self._values + 0.00001).astype(np.int16)
             else:
