FIX

jklymak · jklymak · commit e16ab73a1ae4 · 2020-05-06T22:03:50.000-07:00
diff --git a/lib/matplotlib/gridspec.py b/lib/matplotlib/gridspec.py
@@ -197,6 +197,29 @@ def get_grid_positions(self, fig, raw=False):
         fig_lefts, fig_rights = (left + cell_ws).reshape((-1, 2)).T
         return fig_bottoms, fig_tops, fig_lefts, fig_rights
 
+    @staticmethod
+    def _check_gridspec_exists(figure, nrows, ncols):
+        """
+        Check if the figure already has a gridspec with these dimensions...
+        """
+        gs = None
+        for ax in figure.get_axes():
+             if hasattr(ax, 'get_subplotspec'):
+                 ggs = ax.get_subplotspec().get_gridspec()
+                 if hasattr(ggs, 'get_topmost_subplotspec'):
+                     # This is needed for colorbar gridspec layouts.
+                     # This is probably OK becase this whole logic tree
+                     # is for when the user is doing simple things with the
+                     # add_subplot command.  Complicated stuff, the proper
+                     # gridspec is passed in...
+                     ggs = ggs.get_topmost_subplotspec().get_gridspec()
+
+                 (nrow, ncol) = ggs.get_geometry()
+                 if nrow == nrows and ncol == ncols:
+                     gs = ggs
+        return gs
+
+
     def __getitem__(self, key):
         """Create and return a `.SubplotSpec` instance."""
         nrows, ncols = self.get_geometry()
@@ -558,21 +581,11 @@ def _from_subplot_args(figure, args):
         else:
             raise TypeError(f"subplot() takes 1 or 3 positional arguments but "
                             f"{len(args)} were given")
-        for ax in figure.get_axes():
-             if hasattr(ax, 'get_subplotspec'):
-                 gs = ax.get_subplotspec().get_gridspec()
-                 if hasattr(gs, 'get_topmost_subplotspec'):
-                     # This is needed for colorbar gridspec layouts.
-                     # This is probably OK becase this whole logic tree
-                     # is for when the user is doing simple things with the
-                     # add_subplot command.  Complicated stuff, the proper
-                     # gridspec is passed in...
-                     gs = gs.get_topmost_subplotspec().get_gridspec()
 
-                 (nrow, ncol) = gs.get_geometry()
-                 if nrow == rows and ncol == cols:
-                     return gs[ii-1:jj]
-        return GridSpec(rows, cols, figure=figure)[ii-1:jj]
+        gs = GridSpec._check_gridspec_exists(figure, rows, cols)
+        if gs is None:
+            gs = GridSpec(rows, cols, figure=figure)
+        return gs[ii-1:jj]
 
 
     # num2 is a property only to handle the case where it is None and someone
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -1237,24 +1237,10 @@ def subplot2grid(shape, loc, rowspan=1, colspan=1, fig=None, **kwargs):
     if fig is None:
         fig = gcf()
 
-    s1, s2 = shape
-    gs = None
-    for ax in fig.get_axes():
-         if hasattr(ax, 'get_subplotspec'):
-             ggs = ax.get_subplotspec().get_gridspec()
-             if hasattr(ggs, 'get_topmost_subplotspec'):
-                 # This is needed for colorbar gridspec layouts.
-                 # This is probably OK becase this whole logic tree
-                 # is for when the user is doing simple things with the
-                 # add_subplot command.  Complicated stuff, the proper
-                 # gridspec is passed in...
-                 ggs = ggs.get_topmost_subplotspec().get_gridspec()
-
-             (nrow, ncol) = ggs.get_geometry()
-             if nrow == s1 and ncol == s2:
-                 gs = ggs
+    rows, cols = shape
+    gs = GridSpec._check_gridspec_exists(fig, rows, cols)
     if gs is None:
-        gs = GridSpec(s1, s2, figure=fig)
+        gs = GridSpec(rows, cols, figure=fig)
 
     subplotspec = gs.new_subplotspec(loc, rowspan=rowspan, colspan=colspan)
     ax = fig.add_subplot(subplotspec, **kwargs)
