FIX: simplify compress

jklymak · jklymak · commit b6b8f100ec6e · 2022-03-29T11:15:03.000+02:00
diff --git a/lib/matplotlib/_constrained_layout.py b/lib/matplotlib/_constrained_layout.py
@@ -255,36 +255,41 @@ def check_no_collapsed_axes(layoutgrids, fig):
 
 
 def compress_fixed_aspect(layoutgrids, fig):
-    extraw = dict()
-    extrah = dict()
+    gs = None
     for ax in fig.axes:
         if hasattr(ax, 'get_subplotspec'):
             actual = ax.get_position(original=False)
             ax.apply_aspect()
             sub = ax.get_subplotspec()
-            gs = sub.get_gridspec()
-            if gs not in extraw:
-                extraw[gs] = np.zeros(gs.ncols)
-                extrah[gs] = np.zeros(gs.nrows)
+            _gs = sub.get_gridspec()
+            if gs is None:
+                gs = _gs
+                extraw = np.zeros(gs.ncols)
+                extrah = np.zeros(gs.nrows)
+            elif _gs != gs:
+                raise ValueError('Cannot do compressed layout if axes are not',
+                                 'all from the same gridspec')
             orig = ax.get_position(original=True)
             actual = ax.get_position(original=False)
             dw = orig.width - actual.width
             if dw > 0:
-                extraw[gs][sub.colspan] = np.maximum(extraw[gs][sub.colspan],
-                                                     dw)
+                extraw[sub.colspan] = np.maximum(extraw[sub.colspan], dw)
             dh = orig.height - actual.height
             if dh > 0:
-                extrah[gs][sub.rowspan] = np.maximum(extrah[gs][sub.rowspan],
-                                                     dh)
+                extrah[sub.rowspan] = np.maximum(extrah[sub.rowspan], dh)
 
-    w = np.sum(extraw[gs]) / 2
-    layoutgrids[fig].edit_margin_min('left', w)
-    layoutgrids[fig].edit_margin_min('right', w)
+    if gs is not None:
+        w = np.sum(extraw) / 2
+        layoutgrids[fig].edit_margin_min('left', w)
+        layoutgrids[fig].edit_margin_min('right', w)
 
-    h = np.sum(extrah[gs]) / 2
-    layoutgrids[fig].edit_margin_min('top', h)
-    layoutgrids[fig].edit_margin_min('bottom', h)
-    return layoutgrids
+        h = np.sum(extrah) / 2
+        layoutgrids[fig].edit_margin_min('top', h)
+        layoutgrids[fig].edit_margin_min('bottom', h)
+        return layoutgrids
+    else:
+        raise ValueError('Cannot do compressed layout if no axes ',
+                         'are part of a gridspec.')
 
 
 def get_margin_from_padding(obj, *, w_pad=0, h_pad=0,
