Check gridspecness of colorbars on the right figure.

anntzer · anntzer · commit 48aa8949ec62 · 2023-05-14T23:19:23.000+02:00
... which is the figure on which the colorbar is actually added, and not
necessarily the figure on which colorbar() is called.
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -1271,32 +1271,36 @@ def colorbar(
         if ax is None:
             ax = getattr(mappable, "axes", None)
 
-        if (self.get_layout_engine() is not None and
-                not self.get_layout_engine().colorbar_gridspec):
-            use_gridspec = False
         if cax is None:
             if ax is None:
                 raise ValueError(
                     'Unable to determine Axes to steal space for Colorbar. '
                     'Either provide the *cax* argument to use as the Axes for '
                     'the Colorbar, provide the *ax* argument to steal space '
                     'from it, or add *mappable* to an Axes.')
-            current_ax = self.gca()
+            fig = (  # Figure of first axes; logic copied from make_axes.
+                [*ax.flat] if isinstance(ax, np.ndarray)
+                else [*ax] if np.iterable(ax)
+                else [ax])[0].figure
+            current_ax = fig.gca()
+            if (fig.get_layout_engine() is not None and
+                    not fig.get_layout_engine().colorbar_gridspec):
+                use_gridspec = False
             if (use_gridspec
                     and isinstance(ax, mpl.axes._base._AxesBase)
                     and ax.get_subplotspec()):
                 cax, kwargs = cbar.make_axes_gridspec(ax, **kwargs)
             else:
                 cax, kwargs = cbar.make_axes(ax, **kwargs)
             # make_axes calls add_{axes,subplot} which changes gca; undo that.
-            self.sca(current_ax)
+            fig.sca(current_ax)
             cax.grid(visible=False, which='both', axis='both')
 
         NON_COLORBAR_KEYS = [  # remove kws that cannot be passed to Colorbar
             'fraction', 'pad', 'shrink', 'aspect', 'anchor', 'panchor']
         cb = cbar.Colorbar(cax, mappable, **{
             k: v for k, v in kwargs.items() if k not in NON_COLORBAR_KEYS})
-        self.stale = True
+        cax.figure.stale = True
         return cb
 
     def subplots_adjust(self, left=None, bottom=None, right=None, top=None,
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -1217,3 +1217,16 @@ def test_colorbar_axes_parmeters():
     fig.colorbar(im, ax=(ax[0], ax[1]))
     fig.colorbar(im, ax={i: _ax for i, _ax in enumerate(ax)}.values())
     fig.draw_without_rendering()
+
+
+def test_colorbar_wrong_figure():
+    # If we decide in the future to disallow calling colorbar() on the "wrong" figure,
+    # just delete this test.
+    fig_tl = plt.figure(layout="tight")
+    fig_cl = plt.figure(layout="constrained")
+    im = fig_cl.add_subplot().imshow([[0, 1]])
+    # Make sure this doesn't try to setup a gridspec-controlled colorbar on fig_cl,
+    # which would crash CL.
+    fig_tl.colorbar(im)
+    fig_tl.draw_without_rendering()
+    fig_cl.draw_without_rendering()
