Check if the mappable is in a different Figure than the one fig.color… (#27458)

zoehcycy · web-flow · commit e5562bb7453c · 2023-12-20T14:13:51.000-05:00
* Check if the mappable is in a different Figure than the one fig.colorbar is being called on.
Ignore this check if the mappable has no host Figure.
Warn in case of a mismatch.

* Add figure repr in warning message.

* Use _api.warn_external instead of _log.warning. Only warn when the mismatch is inferred.

* Delete test_colorbar::test_colorbar_wrong_figure.

* Warn in all mismatches, inferred or not.

* Revert deletion of test_colorbar_wrong_figure, with warning test.

* Linting fixes

* Catch newly added warning raised in other test cases.

* Catch warning in existing test cases.

* edge case: subfigure. only warn when top level artist don't match.

* flake8 code style check fix

* code style fix.

* subfigure test cases revert: no warning.

* Fix edge case: axes.remove(); passing correct  pcolormesh object in galleries.

* remove redundant lines

* code stylistic fix
diff --git a/galleries/users_explain/axes/colorbar_placement.py b/galleries/users_explain/axes/colorbar_placement.py
@@ -53,7 +53,7 @@
 fig, axs = plt.subplots(2, 1, figsize=(4, 5), sharex=True)
 X = np.random.randn(20, 20)
 axs[0].plot(np.sum(X, axis=0))
-axs[1].pcolormesh(X)
+pcm = axs[1].pcolormesh(X)
 fig.colorbar(pcm, ax=axs[1], shrink=0.6)
 
 # %%
@@ -63,7 +63,7 @@
 
 fig, axs = plt.subplots(2, 1, figsize=(4, 5), sharex=True, layout='constrained')
 axs[0].plot(np.sum(X, axis=0))
-axs[1].pcolormesh(X)
+pcm = axs[1].pcolormesh(X)
 fig.colorbar(pcm, ax=axs[1], shrink=0.6)
 
 # %%
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -1242,6 +1242,19 @@ def colorbar(
             fig.sca(current_ax)
             cax.grid(visible=False, which='both', axis='both')
 
+        if hasattr(mappable, "figure") and mappable.figure is not None:
+            # Get top level artists
+            mappable_host_fig = mappable.figure
+            if isinstance(mappable_host_fig, mpl.figure.SubFigure):
+                mappable_host_fig = mappable_host_fig.figure
+            # Warn in case of mismatch
+            if mappable_host_fig is not self.figure:
+                _api.warn_external(
+                        f'Adding colorbar to a different Figure '
+                        f'{repr(mappable.figure)} than '
+                        f'{repr(self.figure)} which '
+                        f'fig.colorbar is called on.')
+
         NON_COLORBAR_KEYS = [  # remove kws that cannot be passed to Colorbar
             'fraction', 'pad', 'shrink', 'aspect', 'anchor', 'panchor']
         cb = cbar.Colorbar(cax, mappable, **{
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -1228,7 +1228,8 @@ def test_colorbar_wrong_figure():
     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)
+    with pytest.warns(UserWarning, match="different Figure"):
+        fig_tl.colorbar(im)
     fig_tl.draw_without_rendering()
     fig_cl.draw_without_rendering()
 
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
@@ -1358,7 +1358,6 @@ def test_subfigure_double():
         ax.set_xlabel('x-label', fontsize=fontsize)
         ax.set_ylabel('y-label', fontsize=fontsize)
         ax.set_title('Title', fontsize=fontsize)
-
     subfigsnest[0].colorbar(pc, ax=axsnest0)
 
     subfigsnest[1].suptitle('subfigsnest[1]')
@@ -1665,3 +1664,39 @@ def test_not_visible_figure():
     fig.savefig(buf, format='svg')
     buf.seek(0)
     assert '<g ' not in buf.read()
+
+
+def test_warn_colorbar_mismatch():
+    fig1, ax1 = plt.subplots()
+    fig2, (ax2_1, ax2_2) = plt.subplots(2)
+    im = ax1.imshow([[1, 2], [3, 4]])
+
+    fig1.colorbar(im)  # should not warn
+    with pytest.warns(UserWarning, match="different Figure"):
+        fig2.colorbar(im)
+    # warn mismatch even when the host figure is not inferred
+    with pytest.warns(UserWarning, match="different Figure"):
+        fig2.colorbar(im, ax=ax1)
+    with pytest.warns(UserWarning, match="different Figure"):
+        fig2.colorbar(im, ax=ax2_1)
+    with pytest.warns(UserWarning, match="different Figure"):
+        fig2.colorbar(im, cax=ax2_2)
+
+    # edge case: only compare top level artist in case of subfigure
+    fig3 = plt.figure()
+    fig4 = plt.figure()
+    subfig3_1 = fig3.subfigures()
+    subfig3_2 = fig3.subfigures()
+    subfig4_1 = fig4.subfigures()
+    ax3_1 = subfig3_1.subplots()
+    ax3_2 = subfig3_1.subplots()
+    ax4_1 = subfig4_1.subplots()
+    im3_1 = ax3_1.imshow([[1, 2], [3, 4]])
+    im3_2 = ax3_2.imshow([[1, 2], [3, 4]])
+    im4_1 = ax4_1.imshow([[1, 2], [3, 4]])
+
+    fig3.colorbar(im3_1)   # should not warn
+    subfig3_1.colorbar(im3_1)   # should not warn
+    subfig3_1.colorbar(im3_2)   # should not warn
+    with pytest.warns(UserWarning, match="different Figure"):
+        subfig3_1.colorbar(im4_1)
