MNT: better arg handling GridSpecFromSubplotSpec

jklymak · jklymak · commit db511e847ce7 · 2023-12-24T09:00:51.000-08:00
diff --git a/lib/matplotlib/gridspec.py b/lib/matplotlib/gridspec.py
@@ -484,7 +484,16 @@ def __init__(self, nrows, ncols,
         """
         self._wspace = wspace
         self._hspace = hspace
-        self._subplot_spec = subplot_spec
+        if hasattr(subplot_spec, 'get_subplotspec'):
+            # user has probably passed an axes instead, but
+            # be forgiving.
+            subplot_spec = subplot_spec.get_subplotspec()
+        if isinstance(subplot_spec, SubplotSpec):
+            self._subplot_spec = subplot_spec
+        else:
+            raise TypeError(
+                            "subplot_spec must be type SubplotSpec, "
+                            "usually from GridSpec, or axes.get_subplotspec.")
         self.figure = self._subplot_spec.get_gridspec().figure
         super().__init__(nrows, ncols,
                          width_ratios=width_ratios,
diff --git a/lib/matplotlib/tests/test_gridspec.py b/lib/matplotlib/tests/test_gridspec.py
@@ -1,4 +1,5 @@
 import matplotlib.gridspec as gridspec
+import matplotlib.pyplot as plt
 import pytest
 
 
@@ -35,3 +36,17 @@ def test_repr():
                            width_ratios=(1, 3))
     assert repr(ss) == \
         "GridSpec(2, 2, height_ratios=(3, 1), width_ratios=(1, 3))"
+
+
+def test_subplotspec_args():
+    fig, axs = plt.subplots(1, 2)
+    # should work:
+    gs = gridspec.GridSpecFromSubplotSpec(2, 1,
+                                          subplot_spec=axs[0].get_subplotspec())
+    assert gs.get_topmost_subplotspec() == axs[0].get_subplotspec()
+    # this is a mistake, and not what the type hints give, but we allow:
+    gs = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=axs[0])
+    assert gs.get_topmost_subplotspec() == axs[0].get_subplotspec()
+    # anything else here should type error:
+    with pytest.raises(TypeError, match="subplot_spec must be type SubplotSpec"):
+        gs = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=axs)
