Merge pull request #21751 from jklymak/fix-manual-colorbar-tight

dstansby · web-flow · commit 25a54a4fef67 · 2021-11-28T13:54:03.000Z
FIX: manual colorbars and tight layout
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -280,6 +280,8 @@ def get_subplotspec(self):
         # make tight_layout happy..
         ss = getattr(self._cbar.ax, 'get_subplotspec', None)
         if ss is None:
+            if self._orig_locator is None:
+                return None
             ss = self._orig_locator.get_subplotspec()
         else:
             ss = ss()
diff --git a/lib/matplotlib/tests/test_tightlayout.py b/lib/matplotlib/tests/test_tightlayout.py
@@ -329,3 +329,16 @@ def __init__(self, *args, **kwargs):
     monkeypatch.setattr(mpl.backend_bases.RendererBase, "__init__", __init__)
     fig, ax = plt.subplots()
     fig.tight_layout()
+
+
+def test_manual_colorbar():
+    # This should warn, but not raise
+    fig, axes = plt.subplots(1, 2)
+    pts = axes[1].scatter([0, 1], [0, 1], c=[1, 5])
+    ax_rect = axes[1].get_position()
+    cax = fig.add_axes(
+        [ax_rect.x1 + 0.005, ax_rect.y0, 0.015, ax_rect.height]
+    )
+    fig.colorbar(pts, cax=cax)
+    with pytest.warns(UserWarning, match="This figure includes Axes"):
+        fig.tight_layout()
diff --git a/lib/matplotlib/tight_layout.py b/lib/matplotlib/tight_layout.py
@@ -228,13 +228,14 @@ def get_subplotspec_list(axes_list, grid_spec=None):
 
         if hasattr(axes_or_locator, "get_subplotspec"):
             subplotspec = axes_or_locator.get_subplotspec()
-            subplotspec = subplotspec.get_topmost_subplotspec()
-            gs = subplotspec.get_gridspec()
-            if grid_spec is not None:
-                if gs != grid_spec:
+            if subplotspec is not None:
+                subplotspec = subplotspec.get_topmost_subplotspec()
+                gs = subplotspec.get_gridspec()
+                if grid_spec is not None:
+                    if gs != grid_spec:
+                        subplotspec = None
+                elif gs.locally_modified_subplot_params():
                     subplotspec = None
-            elif gs.locally_modified_subplot_params():
-                subplotspec = None
         else:
             subplotspec = None
 
