Check gridspecness of colorbars on the right figure.

... which is the figure on which the colorbar is actually added, and not
necessarily the figure on which colorbar() is called.

Author: anntzer

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,

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()