ENH: check if colorbar added

jklymak · jklymak · commit e0bb29120bdf · 2021-12-24T10:59:21.000+01:00
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -2336,6 +2336,23 @@ def __init__(self,
         # list of child gridspecs for this figure
         self._gridspecs = []
 
+    def _check_layout_engines_compat(self, old, new):
+        """
+        Helper for set_layout engine
+
+        If the figure has used the old engine and added a colorbar then the
+        value of colorbar_gridspec must be the same on the new engine.
+        """
+        if old is None or old.colorbar_gridspec == new.colorbar_gridspec:
+            return True
+        # colorbar layout different, so check if any colorbars are on the
+        # figure...
+        for ax in self.axes:
+            if hasattr(ax, '_colorbar'):
+                # colorbars list themselvs as a colorbar.
+                return False
+        return True
+
     def set_layout_engine(self, layout=None, **kwargs):
         """
         Set the layout engine for this figure.
@@ -2359,14 +2376,21 @@ def set_layout_engine(self, layout=None, **kwargs):
                 self._layout_engine = None
                 return
         if layout == 'tight':
-            self._layout_engine = TightLayoutEngine(**kwargs)
+            new_layout_engine = TightLayoutEngine(**kwargs)
         elif layout == 'constrained':
-            self._layout_engine = ConstrainedLayoutEngine(**kwargs)
+            new_layout_engine = ConstrainedLayoutEngine(**kwargs)
         elif isinstance(layout, LayoutEngine):
-            self._layout_engine = layout
+            new_layout_engine = layout
         else:
             raise ValueError(f"Invalid value for 'layout': {layout!r}")
-        self._layout_engine.set_figure(self)
+
+        if self._check_layout_engines_compat(self._layout_engine,
+                                             new_layout_engine):
+            self._layout_engine = new_layout_engine
+        else:
+            _api.warn_external('Colorbar layout of new layout engine not '
+                               'compatible with old engine, and a colorbar '
+                               'has been created.  Engine not changed.')
 
     def get_layout_engine(self):
         return self._layout_engine
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
@@ -598,6 +598,30 @@ def test_invalid_layouts():
                        match="Invalid value for 'layout'"):
         Figure(layout='foobar')
 
+    # test that layouts can be swapped if no colorbar:
+    fig, ax = plt.subplots(layout="constrained")
+    fig.set_layout_engine("tight")
+    assert isinstance(fig.get_layout_engine(), TightLayoutEngine)
+    fig.set_layout_engine("constrained")
+    assert isinstance(fig.get_layout_engine(), ConstrainedLayoutEngine)
+
+    # test that layouts cannot be swapped if there is a colorbar:
+    fig, ax = plt.subplots(layout="constrained")
+    pc = ax.pcolormesh(np.random.randn(2, 2))
+    fig.colorbar(pc)
+    with pytest.warns(UserWarning, match='Colorbar layout of new layout'):
+        fig.set_layout_engine("tight")
+    # check that engine didn't change
+    assert isinstance(fig.get_layout_engine(), ConstrainedLayoutEngine)
+
+    fig, ax = plt.subplots(layout="tight")
+    pc = ax.pcolormesh(np.random.randn(2, 2))
+    fig.colorbar(pc)
+    with pytest.warns(UserWarning, match='Colorbar layout of new layout'):
+        fig.set_layout_engine("constrained")
+    # check that engine didn't change
+    assert isinstance(fig.get_layout_engine(), TightLayoutEngine)
+
 
 @check_figures_equal(extensions=["png", "pdf"])
 def test_add_artist(fig_test, fig_ref):
