MNT: better arg handling GridSpecFromSubplotSpec

Author: jklymak

lib/matplotlib/gridspec.py

@@ -484,7 +484,12 @@ def __init__(self, nrows, ncols,
         """
         self._wspace = wspace
         self._hspace = hspace
-        self._subplot_spec = subplot_spec
+        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,

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,15 @@ 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()
+    with pytest.raises(TypeError, match="subplot_spec must be type SubplotSpec"):
+        gs = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=axs[0])
+    with pytest.raises(TypeError, match="subplot_spec must be type SubplotSpec"):
+        gs = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=axs)