enable sca on subfigure axes

rcomer · rcomer · commit 101a2d5eb8ed · 2022-05-09T11:51:30.000+01:00
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -59,7 +59,7 @@
 from matplotlib import cbook
 from matplotlib import _docstring
 from matplotlib.backend_bases import FigureCanvasBase, MouseButton
-from matplotlib.figure import Figure, figaspect
+from matplotlib.figure import Figure, FigureBase, figaspect
 from matplotlib.gridspec import GridSpec, SubplotSpec
 from matplotlib import rcParams, rcParamsDefault, get_backend, rcParamsOrig
 from matplotlib.rcsetup import interactive_bk as _interactive_bk
@@ -690,7 +690,7 @@ def figure(num=None,  # autoincrement if None, else integer from 1-N
 
     Parameters
     ----------
-    num : int or str or `.Figure`, optional
+    num : int or str or `.Figure` or `.SubFigure`, optional
         A unique identifier for the figure.
 
         If a figure with that identifier already exists, this figure is made
@@ -702,7 +702,8 @@ def figure(num=None,  # autoincrement if None, else integer from 1-N
         will be used for the ``Figure.number`` attribute, otherwise, an
         auto-generated integer value is used (starting at 1 and incremented
         for each new figure). If *num* is a string, the figure label and the
-        window title is set to this value.
+        window title is set to this value.  If num is a ``SubFigure``, its
+        parent ``Figure`` is activated.
 
     figsize : (float, float), default: :rc:`figure.figsize`
         Width, height in inches.
@@ -753,11 +754,11 @@ def figure(num=None,  # autoincrement if None, else integer from 1-N
     `~matplotlib.rcParams` defines the default values, which can be modified
     in the matplotlibrc file.
     """
-    if isinstance(num, Figure):
+    if isinstance(num, FigureBase):
         if num.canvas.manager is None:
             raise ValueError("The passed figure is not managed by pyplot")
         _pylab_helpers.Gcf.set_active(num.canvas.manager)
-        return num
+        return num.figure
 
     allnums = get_fignums()
     next_num = max(allnums) + 1 if allnums else 1
diff --git a/lib/matplotlib/tests/test_pyplot.py b/lib/matplotlib/tests/test_pyplot.py
@@ -343,3 +343,27 @@ def test_fallback_position():
     axtest = plt.axes([0.2, 0.2, 0.5, 0.5], position=[0.1, 0.1, 0.8, 0.8])
     np.testing.assert_allclose(axtest.bbox.get_points(),
                                axref.bbox.get_points())
+
+
+def test_set_current_figure_via_subfigure():
+    fig1 = plt.figure()
+    subfigs = fig1.subfigures(2)
+
+    plt.figure()
+    assert plt.gcf() != fig1
+
+    current = plt.figure(subfigs[1])
+    assert plt.gcf() == fig1
+    assert current == fig1
+
+
+def test_set_current_axes_on_subfigure():
+    fig = plt.figure()
+    subfigs = fig.subfigures(2)
+
+    ax = subfigs[0].subplots(1, squeeze=True)
+    subfigs[1].subplots(1, squeeze=True)
+
+    assert plt.gca() != ax
+    plt.sca(ax)
+    assert plt.gca() == ax
