ENH: reuse gridspec if possible

jklymak · jklymak · commit 602dfbe3b753 · 2020-05-06T16:53:18.000-07:00
diff --git a/lib/matplotlib/gridspec.py b/lib/matplotlib/gridspec.py
@@ -536,7 +536,8 @@ def _from_subplot_args(figure, args):
                     raise ValueError("Single argument to subplot must be a "
                                      "3-digit integer") from err
                 # num - 1 for converting from MATLAB to python indexing
-                return GridSpec(rows, cols, figure=figure)[num - 1]
+                ii = num
+                jj = num
         elif len(args) == 3:
             rows, cols, num = args
             rows = int(rows)
@@ -546,17 +547,33 @@ def _from_subplot_args(figure, args):
             if cols <= 0:
                 raise ValueError(f"Number of columns must be > 0, not {cols}")
             if isinstance(num, tuple) and len(num) == 2:
-                i, j = map(int, num)
-                return GridSpec(rows, cols, figure=figure)[i-1:j]
+                ii, jj = map(int, num)
             else:
                 if num < 1 or num > rows*cols:
                     raise ValueError(
                         f"num must be 1 <= num <= {rows*cols}, not {num}")
                 # num - 1 for converting from MATLAB to python indexing
-                return GridSpec(rows, cols, figure=figure)[int(num) - 1]
+                ii = num
+                jj = num
         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]
+
 
     # num2 is a property only to handle the case where it is None and someone
     # mutates num1.
