Fix colorbar dividers with explicit limits

QuLogic · QuLogic · commit 930bd2f03161 · 2022-08-09T01:19:48.000-04:00
In this case, place dividers at the same place internally, but then
always place a divider at the Axes limits (when there's an extend
triangle there).
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -651,12 +651,31 @@ def _add_solids(self, X, Y, C):
             if not self.drawedges:
                 if len(self._y) >= self.n_rasterize:
                     self.solids.set_rasterized(True)
-        if self.drawedges:
-            start_idx = 0 if self._extend_lower() else 1
-            end_idx = len(X) if self._extend_upper() else -1
-            self.dividers.set_segments(np.dstack([X, Y])[start_idx:end_idx])
-        else:
+        self._update_dividers()
+
+    def _update_dividers(self):
+        if not self.drawedges:
             self.dividers.set_segments([])
+            return
+        # Place all *internal* dividers.
+        if self.orientation == 'vertical':
+            lims = self.ax.get_ylim()
+            bounds = (lims[0] < self._y) & (self._y < lims[1])
+        else:
+            lims = self.ax.get_xlim()
+            bounds = (lims[0] < self._y) & (self._y < lims[1])
+        y = self._y[bounds]
+        # And then add outer dividers if extensions are on.
+        if self._extend_lower():
+            y = np.insert(y, 0, lims[0])
+        if self._extend_upper():
+            y = np.append(y, lims[1])
+        X, Y = np.meshgrid([0, 1], y)
+        if self.orientation == 'vertical':
+            segments = np.dstack([X, Y])
+        else:
+            segments = np.dstack([Y, X])
+        self.dividers.set_segments(segments)
 
     def _add_solids_patches(self, X, Y, C, mappable):
         hatches = mappable.hatches * len(C)  # Have enough hatches.
@@ -761,7 +780,8 @@ def _do_extends(self, ax=None):
                 zorder=np.nextafter(self.ax.patch.zorder, -np.inf))
             self.ax.add_patch(patch)
             self._extend_patches.append(patch)
-        return
+
+        self._update_dividers()
 
     def add_lines(self, *args, **kwargs):
         """
diff --git a/lib/matplotlib/tests/baseline_images/test_colorbar/extend_drawedges.png b/lib/matplotlib/tests/baseline_images/test_colorbar/extend_drawedges.png
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -931,28 +931,51 @@ def test_proportional_colorbars():
             fig.colorbar(CS3, spacing=spacings[j], ax=axs[i, j])
 
 
-@pytest.mark.parametrize("extend, coloroffset, res", [
-    ('both', 1, [np.array([[0., 0.], [0., 1.]]),
-                 np.array([[1., 0.], [1., 1.]]),
-                 np.array([[2., 0.], [2., 1.]])]),
-    ('min', 0, [np.array([[0., 0.], [0., 1.]]),
-                np.array([[1., 0.], [1., 1.]])]),
-    ('max', 0, [np.array([[1., 0.], [1., 1.]]),
-                np.array([[2., 0.], [2., 1.]])]),
-    ('neither', -1, [np.array([[1., 0.], [1., 1.]])])
-    ])
-def test_colorbar_extend_drawedges(extend, coloroffset, res):
-    cmap = plt.get_cmap("viridis")
-    bounds = np.arange(3)
-    nb_colors = len(bounds) + coloroffset
-    colors = cmap(np.linspace(100, 255, nb_colors).astype(int))
-    cmap, norm = mcolors.from_levels_and_colors(bounds, colors, extend=extend)
-
-    plt.figure(figsize=(5, 1))
-    ax = plt.subplot(111)
-    cbar = Colorbar(ax, cmap=cmap, norm=norm, orientation='horizontal',
-                    drawedges=True)
-    assert np.all(np.equal(cbar.dividers.get_segments(), res))
+@image_comparison(['extend_drawedges.png'], remove_text=True, style='mpl20')
+def test_colorbar_extend_drawedges():
+    params = [
+        ('both', 1, [[[1.1, 0], [1.1, 1]],
+                     [[2, 0], [2, 1]],
+                     [[2.9, 0], [2.9, 1]]]),
+        ('min', 0, [[[1.1, 0], [1.1, 1]],
+                    [[2, 0], [2, 1]]]),
+        ('max', 0, [[[2, 0], [2, 1]],
+                    [[2.9, 0], [2.9, 1]]]),
+        ('neither', -1, [[[2, 0], [2, 1]]]),
+    ]
+
+    plt.rcParams['axes.linewidth'] = 2
+
+    fig = plt.figure(figsize=(10, 4))
+    subfigs = fig.subfigures(1, 2)
+
+    for orientation, subfig in zip(['horizontal', 'vertical'], subfigs):
+        if orientation == 'horizontal':
+            axs = subfig.subplots(4, 1)
+        else:
+            axs = subfig.subplots(1, 4)
+        fig.subplots_adjust(left=0.05, bottom=0.05, right=0.95, top=0.95)
+
+        for ax, (extend, coloroffset, res) in zip(axs, params):
+            cmap = plt.get_cmap("viridis")
+            bounds = np.arange(5)
+            nb_colors = len(bounds) + coloroffset
+            colors = cmap(np.linspace(100, 255, nb_colors).astype(int))
+            cmap, norm = mcolors.from_levels_and_colors(bounds, colors,
+                                                        extend=extend)
+
+            cbar = Colorbar(ax, cmap=cmap, norm=norm, orientation=orientation,
+                            drawedges=True)
+            # Set limits such that only two colours are visible, and the
+            # dividers would be outside the Axes, to ensure that a) they are
+            # not drawn outside, and b) a divider still appears between the
+            # main colour and the extension.
+            if orientation == 'horizontal':
+                ax.set_xlim(1.1, 2.9)
+            else:
+                ax.set_ylim(1.1, 2.9)
+                res = np.array(res)[:, :, [1, 0]]
+            np.testing.assert_array_equal(cbar.dividers.get_segments(), res)
 
 
 def test_negative_boundarynorm():
