Fix removal of colorbars on nested subgridspecs.

anntzer · anntzer · commit 78791894c96e · 2023-11-22T14:45:04.000+01:00
Indiscriminately going back up to the topmost gridspec is wrong in the
presence of nested gridspecs.

We may want to add an API to just go up one level in gridspecs (i.e.
public access to _subplot_spec) but I'd like to first consider whether
we could e.g. merge GridSpecFromSubplotSpec into GridSpec with the only
difference that the former's get_parent() (or another similarly named
new method) returns another gridspec whereas the latter's returns a
figure.
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -1035,14 +1035,11 @@ def remove(self):
         except AttributeError:
             return
         try:
-            gs = ax.get_subplotspec().get_gridspec()
-            subplotspec = gs.get_topmost_subplotspec()
-        except AttributeError:
-            # use_gridspec was False
+            subplotspec = self.ax.get_subplotspec().get_gridspec()._subplot_spec
+        except AttributeError:  # use_gridspec was False
             pos = ax.get_position(original=True)
             ax._set_position(pos)
-        else:
-            # use_gridspec was True
+        else:  # use_gridspec was True
             ax.set_subplotspec(subplotspec)
 
     def _process_values(self):
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -279,13 +279,16 @@ def test_colorbar_single_scatter():
     plt.colorbar(cs)
 
 
-@pytest.mark.parametrize('use_gridspec', [False, True],
-                         ids=['no gridspec', 'with gridspec'])
-def test_remove_from_figure(use_gridspec):
-    """
-    Test `remove` with the specified ``use_gridspec`` setting
-    """
-    fig, ax = plt.subplots()
+@pytest.mark.parametrize('use_gridspec', [True, False])
+@pytest.mark.parametrize('nested_gridspecs', [True, False])
+def test_remove_from_figure(nested_gridspecs, use_gridspec):
+    """Test `remove` with the specified ``use_gridspec`` setting."""
+    fig = plt.figure()
+    if nested_gridspecs:
+        gs = fig.add_gridspec(2, 2)[1, 1].subgridspec(2, 2)
+        ax = fig.add_subplot(gs[1, 1])
+    else:
+        ax = fig.add_subplot()
     sc = ax.scatter([1, 2], [3, 4])
     sc.set_array(np.array([5, 6]))
     pre_position = ax.get_position()
@@ -298,9 +301,7 @@ def test_remove_from_figure(use_gridspec):
 
 
 def test_remove_from_figure_cl():
-    """
-    Test `remove` with constrained_layout
-    """
+    """Test `remove` with constrained_layout."""
     fig, ax = plt.subplots(constrained_layout=True)
     sc = ax.scatter([1, 2], [3, 4])
     sc.set_array(np.array([5, 6]))
