FIX: removing colorbar's axes also removes colorbar

rcomer · rcomer · commit 714ea734f51d · 2026-04-24T10:07:03.000+01:00
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -11,6 +11,7 @@
    End-users most likely won't need to directly use this module's API.
 """
 
+import functools
 import logging
 
 import numpy as np
@@ -194,6 +195,17 @@ def get_subplotspec(self):
             or getattr(self._orig_locator, "get_subplotspec", lambda: None)())
 
 
+def _remove_cbar_axes(ax, cbar):
+    """
+    Replacement remove method for a colorbar's axes, so that the colorbar is
+    properly removed.
+
+    Note we define this at the module level to preserve pickling. A lambda or
+    local def within the Colorbar.__init__ method will not work.
+    """
+    cbar.remove()
+
+
 @_docstring.interpd
 class Colorbar:
     r"""
@@ -427,6 +439,14 @@ def __init__(
         self._extend_cid2 = self.ax.callbacks.connect(
             "ylim_changed", self._do_extends)
 
+        # Ensure proper cleanup when `cbar.ax.remove()` is called.  We ensure
+        # this by overriding the Axes' remove method, so that `cbar.ax.remove()`
+        # actually calls `cbar.remove()`.  In turn, we store the original Axes'
+        # remove method in `_ax_remove`, which `cbar.remove()` will eventually
+        # call to clean up the Axes itself.
+        self._ax_remove = self.ax._remove_method
+        self.ax._remove_method = functools.partial(_remove_cbar_axes, cbar=self)
+
     @property
     def long_axis(self):
         """Axis that has decorations (ticks, etc) on it."""
@@ -1032,7 +1052,7 @@ def remove(self):
                 if self.ax in a._colorbars:
                     a._colorbars.remove(self.ax)
 
-        self.ax.remove()
+        self._ax_remove(self.ax)
 
         self.mappable.callbacks.disconnect(self.mappable.colorbar_cid)
         self.mappable.colorbar = None
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -1,4 +1,5 @@
 import platform
+from unittest import mock
 
 import numpy as np
 import pytest
@@ -311,6 +312,18 @@ def test_remove_from_figure_cl():
                                post_position.get_points())
 
 
+def test_axes_remove():
+    """Removing the colorbar's axes should also remove the colorbar"""
+    fig, ax = plt.subplots()
+    sc = ax.scatter([1, 2], [3, 4])
+    cb = fig.colorbar(sc)
+
+    with mock.patch.object(cb, 'remove') as mock_cb_remove:
+        cb.ax.remove()
+
+    mock_cb_remove.assert_called()
+
+
 def test_colorbarbase():
     # smoke test from #3805
     ax = plt.gca()
