MAINT: Adding cla/remove to the ColorbarAxes class

greglucas · greglucas · commit 51db38e7d677 · 2021-05-28T08:21:00.000-06:00
The cla and remove methods need to be applied to both the
inner and outer axes.
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -285,6 +285,16 @@ def _set_inner_bounds(self, bounds):
         self.inner_ax._axes_locator = _TransformedBoundsLocator(
             bounds, self.outer_ax.transAxes)
 
+    def cla(self):
+        """Clear the axes."""
+        self.inner_ax.cla()
+        self.outer_ax.cla()
+
+    def remove(self):
+        """Remove the axes."""
+        self.inner_ax.remove()
+        self.outer_ax.remove()
+
 
 class _ColorbarSpine(mspines.Spine):
     def __init__(self, axes):
@@ -874,8 +884,7 @@ def set_alpha(self, alpha):
 
     def remove(self):
         """Remove this colorbar from the figure."""
-        self.ax.inner_ax.remove()
-        self.ax.outer_ax.remove()
+        self.ax.remove()
 
     def _ticker(self, locator, formatter):
         """
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -771,3 +771,17 @@ def test_inset_colorbar_layout():
     np.testing.assert_allclose(cb.ax.get_position().bounds,
                                [0.87, 0.342, 0.0237, 0.315], atol=0.01)
     assert cb.ax.outer_ax in ax.child_axes
+
+
+@check_figures_equal(extensions=["png"])
+def test_colorbar_reuse_axes(fig_ref, fig_test):
+    ax = fig_ref.add_subplot()
+    pc = ax.imshow(np.arange(100).reshape(10, 10))
+    cb = fig_ref.colorbar(pc)
+
+    ax = fig_test.add_subplot()
+    pc = ax.imshow(np.arange(100).reshape(10, 10))
+    cb = fig_test.colorbar(pc)
+    # Clear and re-use the same colorbar axes
+    cb.ax.cla()
+    cb = fig_test.colorbar(pc, cax=cb.ax)
