From d6b6504fb8d2a7763a562900e5a6561ae1e13915 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Wed, 6 Dec 2023 20:14:37 -0800
Subject: [PATCH 01/16] 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.

---
 lib/matplotlib/figure.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
index 11b42b1e1ac7..b6cd7cc10a51 100644
--- a/lib/matplotlib/figure.py
+++ b/lib/matplotlib/figure.py
@@ -1242,6 +1242,11 @@ def colorbar(
             fig.sca(current_ax)
             cax.grid(visible=False, which='both', axis='both')
 
+        if hasattr(mappable, "figure") and mappable.figure is not self.figure:
+            _log.warning(
+                    'Trying to add colorbar to a Figure different '
+                    'than the one 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, **{

From c7aac68568d147d270c4241539215af3495d7884 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 01:28:50 -0800
Subject: [PATCH 02/16] Add figure repr in warning message.

---
 lib/matplotlib/figure.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
index b6cd7cc10a51..92d3d842d814 100644
--- a/lib/matplotlib/figure.py
+++ b/lib/matplotlib/figure.py
@@ -1244,8 +1244,9 @@ def colorbar(
 
         if hasattr(mappable, "figure") and mappable.figure is not self.figure:
             _log.warning(
-                    'Trying to add colorbar to a Figure different '
-                    'than the one fig.colorbar is called on.')
+                    'Adding colorbar to a different Figure %s'
+                    ' than %s which fig.colorbar is called on.',
+                    repr(mappable.figure), repr(self.figure))
 
         NON_COLORBAR_KEYS = [  # remove kws that cannot be passed to Colorbar
             'fraction', 'pad', 'shrink', 'aspect', 'anchor', 'panchor']

From 9339a19ba9640acdfdcdd6af2288427a43c369f3 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 06:24:48 -0800
Subject: [PATCH 03/16] Use _api.warn_external instead of _log.warning. Only
 warn when the mismatch is inferred.

---
 lib/matplotlib/figure.py            | 10 +++++-----
 lib/matplotlib/tests/test_figure.py | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
index 92d3d842d814..2636eb093bca 100644
--- a/lib/matplotlib/figure.py
+++ b/lib/matplotlib/figure.py
@@ -1213,6 +1213,7 @@ def colorbar(
         therefore, this workaround is not used by default (see issue #1188).
 
         """
+        infer_ax = (ax is None) and (cax is None)
 
         if ax is None:
             ax = getattr(mappable, "axes", None)
@@ -1242,11 +1243,10 @@ def colorbar(
             fig.sca(current_ax)
             cax.grid(visible=False, which='both', axis='both')
 
-        if hasattr(mappable, "figure") and mappable.figure is not self.figure:
-            _log.warning(
-                    'Adding colorbar to a different Figure %s'
-                    ' than %s which fig.colorbar is called on.',
-                    repr(mappable.figure), repr(self.figure))
+        if infer_ax and hasattr(mappable, "figure") and mappable.figure is not self.figure:
+            _api.warn_external(
+                    f'Adding colorbar to a different Figure {repr(mappable.figure)} '
+                    'than {repr(self.figure)} which fig.colorbar is called on.')
 
         NON_COLORBAR_KEYS = [  # remove kws that cannot be passed to Colorbar
             'fraction', 'pad', 'shrink', 'aspect', 'anchor', 'panchor']
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
index 99b2602bc4a7..7cf90e34612f 100644
--- a/lib/matplotlib/tests/test_figure.py
+++ b/lib/matplotlib/tests/test_figure.py
@@ -1665,3 +1665,21 @@ def test_not_visible_figure():
     fig.savefig(buf, format='svg')
     buf.seek(0)
     assert '<g ' not in buf.read()
+
+def test_warn_colorbar_mismatch():
+    fig, ax = plt.subplots()
+    fig2, (axA, axB) = plt.subplots(2)
+    im = ax.imshow([[1, 2], [3, 4]])
+
+    with pytest.warns() as mismatch_warn:
+        fig2.colorbar(im)   # this should warn
+    assert len(mismatch_warn) == 1
+    assert mismatch_warn[0].category is UserWarning
+    assert str(mismatch_warn[0].message).startswith("Adding colorbar to a different Figure")
+
+    with warnings.catch_warnings() as mismatch_warn:
+        warnings.simplefilter("error")
+        fig2.colorbar(im, ax=ax)   # this is weird, but should not warn
+        fig2.colorbar(im, ax=axA)  # this should not warn
+        fig2.colorbar(im, cax=axB) # also should not warn
+    assert not mismatch_warn

From ba97ec53a57dd16d0178ed5b0209314cde4471e1 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 14:43:35 -0800
Subject: [PATCH 04/16] Delete test_colorbar::test_colorbar_wrong_figure.

---
 lib/matplotlib/tests/test_colorbar.py | 13 -------------
 lib/matplotlib/tests/test_figure.py   |  5 +----
 2 files changed, 1 insertion(+), 17 deletions(-)

diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
index 509d08dae183..99f6dc613700 100644
--- a/lib/matplotlib/tests/test_colorbar.py
+++ b/lib/matplotlib/tests/test_colorbar.py
@@ -1220,19 +1220,6 @@ def test_colorbar_axes_parmeters():
     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()
-
-
 def test_colorbar_format_string_and_old():
     plt.imshow([[0, 1]])
     cb = plt.colorbar(format="{x}%")
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
index 7cf90e34612f..418b31f25917 100644
--- a/lib/matplotlib/tests/test_figure.py
+++ b/lib/matplotlib/tests/test_figure.py
@@ -1671,11 +1671,8 @@ def test_warn_colorbar_mismatch():
     fig2, (axA, axB) = plt.subplots(2)
     im = ax.imshow([[1, 2], [3, 4]])
 
-    with pytest.warns() as mismatch_warn:
+    with pytest.warns(UserWarning, match = "different Figure"):
         fig2.colorbar(im)   # this should warn
-    assert len(mismatch_warn) == 1
-    assert mismatch_warn[0].category is UserWarning
-    assert str(mismatch_warn[0].message).startswith("Adding colorbar to a different Figure")
 
     with warnings.catch_warnings() as mismatch_warn:
         warnings.simplefilter("error")

From 742cd00b8f1877d590979342737dc7c635b4c5ce Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 16:12:22 -0800
Subject: [PATCH 05/16] Warn in all mismatches, inferred or not.

---
 lib/matplotlib/figure.py            | 10 ++++++----
 lib/matplotlib/tests/test_figure.py | 19 ++++++++++---------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
index 2636eb093bca..195f3407781c 100644
--- a/lib/matplotlib/figure.py
+++ b/lib/matplotlib/figure.py
@@ -1213,7 +1213,6 @@ def colorbar(
         therefore, this workaround is not used by default (see issue #1188).
 
         """
-        infer_ax = (ax is None) and (cax is None)
 
         if ax is None:
             ax = getattr(mappable, "axes", None)
@@ -1243,10 +1242,13 @@ def colorbar(
             fig.sca(current_ax)
             cax.grid(visible=False, which='both', axis='both')
 
-        if infer_ax and hasattr(mappable, "figure") and mappable.figure is not self.figure:
+
+        if hasattr(mappable, "figure") and mappable.figure is not self.figure:
             _api.warn_external(
-                    f'Adding colorbar to a different Figure {repr(mappable.figure)} '
-                    'than {repr(self.figure)} which fig.colorbar is called on.')
+                    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']
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
index 418b31f25917..4a097ba2314d 100644
--- a/lib/matplotlib/tests/test_figure.py
+++ b/lib/matplotlib/tests/test_figure.py
@@ -1666,17 +1666,18 @@ def test_not_visible_figure():
     buf.seek(0)
     assert '<g ' not in buf.read()
 
+
 def test_warn_colorbar_mismatch():
     fig, ax = plt.subplots()
     fig2, (axA, axB) = plt.subplots(2)
     im = ax.imshow([[1, 2], [3, 4]])
 
-    with pytest.warns(UserWarning, match = "different Figure"):
-        fig2.colorbar(im)   # this should warn
-
-    with warnings.catch_warnings() as mismatch_warn:
-        warnings.simplefilter("error")
-        fig2.colorbar(im, ax=ax)   # this is weird, but should not warn
-        fig2.colorbar(im, ax=axA)  # this should not warn
-        fig2.colorbar(im, cax=axB) # also should not warn
-    assert not mismatch_warn
+    with pytest.warns(UserWarning, match="different Figure"):
+        fig2.colorbar(im)
+    # warn even when the host figure is not inferred
+    with pytest.warns(UserWarning, match="different Figure"):
+        fig2.colorbar(im, ax=ax)
+    with pytest.warns(UserWarning, match="different Figure"):
+        fig2.colorbar(im, ax=axA) 
+    with pytest.warns(UserWarning, match="different Figure"):
+        fig2.colorbar(im, cax=axB)

From 78bce339fc49e8528b6429624c9264ed447c1d89 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 16:31:08 -0800
Subject: [PATCH 06/16] Revert deletion of test_colorbar_wrong_figure, with
 warning test.

---
 lib/matplotlib/tests/test_colorbar.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
index 99f6dc613700..4566ce449aac 100644
--- a/lib/matplotlib/tests/test_colorbar.py
+++ b/lib/matplotlib/tests/test_colorbar.py
@@ -1220,6 +1220,20 @@ def test_colorbar_axes_parmeters():
     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.
+    with pytest.warns(UserWarning, match="different Figure"):
+        fig_tl.colorbar(im)
+    fig_tl.draw_without_rendering()
+    fig_cl.draw_without_rendering()
+
+
 def test_colorbar_format_string_and_old():
     plt.imshow([[0, 1]])
     cb = plt.colorbar(format="{x}%")

From 10422fda025487d93f9fb86a906e7bd5375f8658 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 16:42:33 -0800
Subject: [PATCH 07/16] Linting fixes

---
 lib/matplotlib/figure.py            | 1 -
 lib/matplotlib/tests/test_figure.py | 2 +-
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
index 195f3407781c..21fc641b1fe5 100644
--- a/lib/matplotlib/figure.py
+++ b/lib/matplotlib/figure.py
@@ -1242,7 +1242,6 @@ def colorbar(
             fig.sca(current_ax)
             cax.grid(visible=False, which='both', axis='both')
 
-
         if hasattr(mappable, "figure") and mappable.figure is not self.figure:
             _api.warn_external(
                     f'Adding colorbar to a different Figure '
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
index 4a097ba2314d..adf7269b3174 100644
--- a/lib/matplotlib/tests/test_figure.py
+++ b/lib/matplotlib/tests/test_figure.py
@@ -1678,6 +1678,6 @@ def test_warn_colorbar_mismatch():
     with pytest.warns(UserWarning, match="different Figure"):
         fig2.colorbar(im, ax=ax)
     with pytest.warns(UserWarning, match="different Figure"):
-        fig2.colorbar(im, ax=axA) 
+        fig2.colorbar(im, ax=axA)
     with pytest.warns(UserWarning, match="different Figure"):
         fig2.colorbar(im, cax=axB)

From 3cbc90a34e167e53a6e3aba450e354e17eebc160 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 18:11:55 -0800
Subject: [PATCH 08/16] Catch newly added warning raised in other test cases.

---
 lib/matplotlib/tests/test_colorbar.py |  5 +++--
 lib/matplotlib/tests/test_figure.py   | 15 ++++++++++-----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
index 4566ce449aac..e44ce06aa13b 100644
--- a/lib/matplotlib/tests/test_colorbar.py
+++ b/lib/matplotlib/tests/test_colorbar.py
@@ -1034,8 +1034,9 @@ def test_colorbar_contourf_extend_patches():
         for ax, (extend, levels, hatches) in zip(axs, params):
             cs = ax.contourf(x, y, z, levels, hatches=hatches,
                              cmap=cmap, extend=extend)
-            subfig.colorbar(cs, ax=ax, orientation=orientation, fraction=0.4,
-                            extendfrac=0.2, aspect=5)
+            with pytest.warns(UserWarning, match="different Figure"):
+                subfig.colorbar(cs, ax=ax, orientation=orientation, fraction=0.4,
+                                extendfrac=0.2, aspect=5)
 
 
 def test_negative_boundarynorm():
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
index adf7269b3174..cdfa12d0c1d3 100644
--- a/lib/matplotlib/tests/test_figure.py
+++ b/lib/matplotlib/tests/test_figure.py
@@ -1280,7 +1280,8 @@ def test_subfigure():
     axs = sub[1].subplots(1, 3)
     for ax in axs.flat:
         pc = ax.pcolormesh(np.random.randn(30, 30), vmin=-2, vmax=2)
-    sub[1].colorbar(pc, ax=axs, location='bottom')
+    with pytest.warns(UserWarning, match="different Figure"):
+        sub[1].colorbar(pc, ax=axs, location='bottom')
     sub[1].suptitle('Right Side')
     sub[1].set_facecolor('white')
 
@@ -1320,7 +1321,8 @@ def test_subfigure_ss():
     axs = sub.subplots(2, 2)
     for ax in axs.flat:
         pc = ax.pcolormesh(np.random.randn(30, 30), vmin=-2, vmax=2)
-    sub.colorbar(pc, ax=axs)
+    with pytest.warns(UserWarning, match="different Figure"):
+        sub.colorbar(pc, ax=axs)
     sub.suptitle('Left Side')
 
     ax = fig.add_subplot(gs[1])
@@ -1358,8 +1360,8 @@ 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)
+    with pytest.warns(UserWarning, match="different Figure"):
+        subfigsnest[0].colorbar(pc, ax=axsnest0)
 
     subfigsnest[1].suptitle('subfigsnest[1]')
     subfigsnest[1].set_facecolor('g')
@@ -1672,9 +1674,12 @@ def test_warn_colorbar_mismatch():
     fig2, (axA, axB) = plt.subplots(2)
     im = ax.imshow([[1, 2], [3, 4]])
 
+    # this should not warn
+    fig.colorbar(im)
+    # warn on an inferred mismatch
     with pytest.warns(UserWarning, match="different Figure"):
         fig2.colorbar(im)
-    # warn even when the host figure is not inferred
+    # warn mismatch even when the host figure is not inferred
     with pytest.warns(UserWarning, match="different Figure"):
         fig2.colorbar(im, ax=ax)
     with pytest.warns(UserWarning, match="different Figure"):

From 2defbed0b4b471bc7909afe6927e0670f793b5cd Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 18:53:48 -0800
Subject: [PATCH 09/16] Catch warning in existing test cases.

---
 lib/matplotlib/tests/test_figure.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
index cdfa12d0c1d3..16e50d9969ab 100644
--- a/lib/matplotlib/tests/test_figure.py
+++ b/lib/matplotlib/tests/test_figure.py
@@ -1273,7 +1273,8 @@ def test_subfigure():
     axs = sub[0].subplots(2, 2)
     for ax in axs.flat:
         pc = ax.pcolormesh(np.random.randn(30, 30), vmin=-2, vmax=2)
-    sub[0].colorbar(pc, ax=axs)
+    with pytest.warns(UserWarning, match="different Figure"):
+        sub[0].colorbar(pc, ax=axs)
     sub[0].suptitle('Left Side')
     sub[0].set_facecolor('white')
 

From c603a4b4a7bf946495aa13bd1c0c7a6c7742dd10 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 21:14:30 -0800
Subject: [PATCH 10/16] edge case: subfigure. only warn when top level artist
 don't match.

---
 lib/matplotlib/figure.py            | 23 +++++++++++------
 lib/matplotlib/tests/test_figure.py | 38 ++++++++++++++++++++---------
 2 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
index 21fc641b1fe5..f6a62a1cedfc 100644
--- a/lib/matplotlib/figure.py
+++ b/lib/matplotlib/figure.py
@@ -1241,13 +1241,22 @@ def colorbar(
             # make_axes calls add_{axes,subplot} which changes gca; undo that.
             fig.sca(current_ax)
             cax.grid(visible=False, which='both', axis='both')
-
-        if hasattr(mappable, "figure") and mappable.figure 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.')
+        
+        if hasattr(mappable, "figure"):
+            # Get top level artists
+            mappable_host_fig = mappable.figure
+            self_host_fig = self.figure
+            if isinstance(mappable_host_fig, mpl.figure.SubFigure):
+                mappable_host_fig = mappable_host_fig.figure
+            if isinstance(self_host_fig, mpl.figure.SubFigure):
+                self_host_fig = self_host_fig.figure
+            # Warn in case of mismatch
+            if mappable_host_fig is not self_host_fig:
+                _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']
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
index 16e50d9969ab..e4ec67c00af7 100644
--- a/lib/matplotlib/tests/test_figure.py
+++ b/lib/matplotlib/tests/test_figure.py
@@ -1281,8 +1281,7 @@ def test_subfigure():
     axs = sub[1].subplots(1, 3)
     for ax in axs.flat:
         pc = ax.pcolormesh(np.random.randn(30, 30), vmin=-2, vmax=2)
-    with pytest.warns(UserWarning, match="different Figure"):
-        sub[1].colorbar(pc, ax=axs, location='bottom')
+    sub[1].colorbar(pc, ax=axs, location='bottom')
     sub[1].suptitle('Right Side')
     sub[1].set_facecolor('white')
 
@@ -1671,19 +1670,36 @@ def test_not_visible_figure():
 
 
 def test_warn_colorbar_mismatch():
-    fig, ax = plt.subplots()
-    fig2, (axA, axB) = plt.subplots(2)
-    im = ax.imshow([[1, 2], [3, 4]])
+    fig1, ax1 = plt.subplots()
+    fig2, (ax2_1, ax2_2) = plt.subplots(2)
+    im = ax1.imshow([[1, 2], [3, 4]])
 
-    # this should not warn
-    fig.colorbar(im)
-    # warn on an inferred mismatch
+    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=ax)
+        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, ax=axA)
+        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"):
-        fig2.colorbar(im, cax=axB)
+        subfig3_1.colorbar(im4_1)

From c1d19ecc4d8e7d50ca9df8bf599f9dd65a5b5698 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 21:19:45 -0800
Subject: [PATCH 11/16] flake8 code style check fix

---
 lib/matplotlib/tests/test_figure.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
index e4ec67c00af7..f0a4b19223b1 100644
--- a/lib/matplotlib/tests/test_figure.py
+++ b/lib/matplotlib/tests/test_figure.py
@@ -1694,9 +1694,9 @@ def test_warn_colorbar_mismatch():
     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]])
+    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

From c14bf186c82ca2cf0c60bf3dea75d35e7d9ad1f8 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 21:23:23 -0800
Subject: [PATCH 12/16] code style fix.

---
 lib/matplotlib/figure.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
index f6a62a1cedfc..8272071d2cb0 100644
--- a/lib/matplotlib/figure.py
+++ b/lib/matplotlib/figure.py
@@ -1241,7 +1241,7 @@ def colorbar(
             # make_axes calls add_{axes,subplot} which changes gca; undo that.
             fig.sca(current_ax)
             cax.grid(visible=False, which='both', axis='both')
-        
+
         if hasattr(mappable, "figure"):
             # Get top level artists
             mappable_host_fig = mappable.figure

From 1c27ec1e3440a27a53fb3225d928a2782cff3ff4 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 21:52:28 -0800
Subject: [PATCH 13/16] subfigure test cases revert: no warning.

---
 lib/matplotlib/tests/test_colorbar.py | 5 ++---
 lib/matplotlib/tests/test_figure.py   | 9 +++------
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
index e44ce06aa13b..4566ce449aac 100644
--- a/lib/matplotlib/tests/test_colorbar.py
+++ b/lib/matplotlib/tests/test_colorbar.py
@@ -1034,9 +1034,8 @@ def test_colorbar_contourf_extend_patches():
         for ax, (extend, levels, hatches) in zip(axs, params):
             cs = ax.contourf(x, y, z, levels, hatches=hatches,
                              cmap=cmap, extend=extend)
-            with pytest.warns(UserWarning, match="different Figure"):
-                subfig.colorbar(cs, ax=ax, orientation=orientation, fraction=0.4,
-                                extendfrac=0.2, aspect=5)
+            subfig.colorbar(cs, ax=ax, orientation=orientation, fraction=0.4,
+                            extendfrac=0.2, aspect=5)
 
 
 def test_negative_boundarynorm():
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
index f0a4b19223b1..0e1706f37b59 100644
--- a/lib/matplotlib/tests/test_figure.py
+++ b/lib/matplotlib/tests/test_figure.py
@@ -1273,8 +1273,7 @@ def test_subfigure():
     axs = sub[0].subplots(2, 2)
     for ax in axs.flat:
         pc = ax.pcolormesh(np.random.randn(30, 30), vmin=-2, vmax=2)
-    with pytest.warns(UserWarning, match="different Figure"):
-        sub[0].colorbar(pc, ax=axs)
+    sub[0].colorbar(pc, ax=axs)
     sub[0].suptitle('Left Side')
     sub[0].set_facecolor('white')
 
@@ -1321,8 +1320,7 @@ def test_subfigure_ss():
     axs = sub.subplots(2, 2)
     for ax in axs.flat:
         pc = ax.pcolormesh(np.random.randn(30, 30), vmin=-2, vmax=2)
-    with pytest.warns(UserWarning, match="different Figure"):
-        sub.colorbar(pc, ax=axs)
+    sub.colorbar(pc, ax=axs)
     sub.suptitle('Left Side')
 
     ax = fig.add_subplot(gs[1])
@@ -1360,8 +1358,7 @@ def test_subfigure_double():
         ax.set_xlabel('x-label', fontsize=fontsize)
         ax.set_ylabel('y-label', fontsize=fontsize)
         ax.set_title('Title', fontsize=fontsize)
-    with pytest.warns(UserWarning, match="different Figure"):
-        subfigsnest[0].colorbar(pc, ax=axsnest0)
+    subfigsnest[0].colorbar(pc, ax=axsnest0)
 
     subfigsnest[1].suptitle('subfigsnest[1]')
     subfigsnest[1].set_facecolor('g')

From 6308cb6f71590aa891f2afe5d701227e65c47437 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Thu, 7 Dec 2023 23:19:01 -0800
Subject: [PATCH 14/16] Fix edge case: axes.remove(); passing correct 
 pcolormesh object in galleries.

---
 galleries/users_explain/axes/colorbar_placement.py | 4 ++--
 lib/matplotlib/figure.py                           | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/galleries/users_explain/axes/colorbar_placement.py b/galleries/users_explain/axes/colorbar_placement.py
index 8dbc2a356cb1..919600b070f3 100644
--- 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
index 8272071d2cb0..ecc1ff0e1ea6 100644
--- a/lib/matplotlib/figure.py
+++ b/lib/matplotlib/figure.py
@@ -1242,7 +1242,7 @@ def colorbar(
             fig.sca(current_ax)
             cax.grid(visible=False, which='both', axis='both')
 
-        if hasattr(mappable, "figure"):
+        if hasattr(mappable, "figure") and mappable.figure is not None:
             # Get top level artists
             mappable_host_fig = mappable.figure
             self_host_fig = self.figure

From ecddddd051cec8132348f1f0f327010e4a0451bc Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Wed, 13 Dec 2023 17:20:39 -0800
Subject: [PATCH 15/16] remove redundant lines

---
 lib/matplotlib/figure.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
index ecc1ff0e1ea6..41b4e89eded4 100644
--- a/lib/matplotlib/figure.py
+++ b/lib/matplotlib/figure.py
@@ -1248,8 +1248,6 @@ def colorbar(
             self_host_fig = self.figure
             if isinstance(mappable_host_fig, mpl.figure.SubFigure):
                 mappable_host_fig = mappable_host_fig.figure
-            if isinstance(self_host_fig, mpl.figure.SubFigure):
-                self_host_fig = self_host_fig.figure
             # Warn in case of mismatch
             if mappable_host_fig is not self_host_fig:
                 _api.warn_external(

From 998a9243964b590a8effa76d09c3e35694e0f097 Mon Sep 17 00:00:00 2001
From: Chaoyi Hu <zoehcycy@gmail.com>
Date: Tue, 19 Dec 2023 21:49:11 -0800
Subject: [PATCH 16/16] code stylistic fix

---
 lib/matplotlib/figure.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
index 41b4e89eded4..dd76c932b01e 100644
--- a/lib/matplotlib/figure.py
+++ b/lib/matplotlib/figure.py
@@ -1245,11 +1245,10 @@ def colorbar(
         if hasattr(mappable, "figure") and mappable.figure is not None:
             # Get top level artists
             mappable_host_fig = mappable.figure
-            self_host_fig = self.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_host_fig:
+            if mappable_host_fig is not self.figure:
                 _api.warn_external(
                         f'Adding colorbar to a different Figure '
                         f'{repr(mappable.figure)} than '