ENH/WIP: make it work recursively

tacaswell · tacaswell · commit a2d4d9a2d783 · 2020-02-28T04:19:35.000-05:00
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -18,7 +18,6 @@
 from matplotlib import rcParams
 from matplotlib import docstring, projections
 from matplotlib import __version__ as _mpl_version
-import matplotlib.gridspec as gridspec
 
 import matplotlib.artist as martist
 from matplotlib.artist import Artist, allow_rasterization
@@ -30,7 +29,7 @@
 
 from matplotlib.axes import Axes, SubplotBase, subplot_class_factory
 from matplotlib.blocking_input import BlockingMouseInput, BlockingKeyMouseInput
-from matplotlib.gridspec import GridSpec
+from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec
 import matplotlib.legend as mlegend
 from matplotlib.patches import Rectangle
 from matplotlib.text import Text
@@ -1617,25 +1616,47 @@ def build_grid(self, layout, subplot_kw=None, gridspec_kw=None):
         """
         subplot_kw = subplot_kw or {}
         gridspec_kw = gridspec_kw or {}
+
+        def _process_layout(layout):
+            unique_ids = set()
+            nested = {}
+            for j, row in enumerate(layout):
+                for k, v in enumerate(row):
+                    if isinstance(v, str):
+                        unique_ids.add(v)
+                    else:
+                        nested[(j, k)] = v
+
+            return unique_ids, nested
+
+        def _do_layout(gs, layout, unique_ids, nested):
+            rows, cols = layout.shape
+            covered = np.zeros((rows, cols), dtype=bool)
+            output = dict()
+            for name in unique_ids:
+                indx = np.stack(np.where(layout == name)).T
+                start_row, start_col = np.min(indx, axis=0)
+                end_row, end_col = np.max(indx, axis=0) + 1
+                slc = (slice(start_row, end_row), slice(start_col, end_col))
+                if not np.all(covered[slc] == False):
+                    raise ValueError
+                covered[slc] = True
+                output[name] = self.add_subplot(gs[slc], **subplot_kw)
+
+            for (j, k), layout in nested.items():
+                layout = np.asarray(layout)
+                rows, cols = layout.shape
+                gs_n = GridSpecFromSubplotSpec(rows, cols, gs[j, k])
+                nested_output = _do_layout(gs_n, layout, *_process_layout(layout))
+
+                output.update(nested_output)
+            return output
+
         layout = np.asarray(layout)
         rows, cols = layout.shape
-        unique_ids = set(layout.ravel())
-
-        gs = gridspec.GridSpec(rows, cols, figure=self, **gridspec_kw)
-
-        covered = np.zeros((rows, cols), dtype=bool)
-        output = dict()
-        for name in unique_ids:
-            indx = np.stack(np.where(layout == name)).T
-            start_row, start_col = np.min(indx, axis=0)
-            end_row, end_col = np.max(indx, axis=0) + 1
-            slc = (slice(start_row, end_row), slice(start_col, end_col))
-            if not np.all(covered[slc] == False):
-                raise ValueError
-            covered[slc] = True
-            output[name] = self.add_subplot(gs[slc], **subplot_kw)
-
-        return output
+        gs = GridSpec(rows, cols, figure=self, **gridspec_kw)
+        return  _do_layout(gs, layout, *_process_layout(layout))
+
 
     def delaxes(self, ax):
         """
diff --git a/lib/matplotlib/tests/test_build_grid.py b/lib/matplotlib/tests/test_build_grid.py
@@ -3,7 +3,7 @@
 
 
 @check_figures_equal()
-def test_label_loc_vertical(fig_test, fig_ref):
+def test_basic(fig_test, fig_ref):
     x = [["A", "A", "B"], ["C", "D", "B"]]
     grid_axes = fig_test.build_grid(x)
 
@@ -22,3 +22,36 @@ def test_label_loc_vertical(fig_test, fig_ref):
 
     axD = fig_ref.add_subplot(gs[1, 1])
     axD.set_title("D")
+
+
+@check_figures_equal(tol=0.005)
+def test_recursive(fig_test, fig_ref):
+    x = [["A", "B", "B"], ["C", "C", "D"]]
+
+    y = [["F"], [x]]
+
+    grid_axes = fig_test.build_grid(y)
+
+    for k, ax in grid_axes.items():
+        ax.set_title(k)
+
+    gs = gridspec.GridSpec(2, 1, figure=fig_ref)
+    axF = fig_ref.add_subplot(gs[0, 0])
+    axF.set_title("F")
+
+    gs_n = gridspec.GridSpecFromSubplotSpec(2, 3, gs[1, 0])
+
+    axA = fig_ref.add_subplot(gs_n[0, 0])
+    axA.set_title("A")
+
+    axB = fig_ref.add_subplot(gs_n[0, 1:])
+    axB.set_title("B")
+
+    axC = fig_ref.add_subplot(gs_n[1, :2])
+    axC.set_title("C")
+
+    axD = fig_ref.add_subplot(gs_n[1, 2])
+    axD.set_title("D")
+
+    fig_test.tight_layout()
+    fig_ref.tight_layout()
