Rationalise artist get_figure; make figure attribute a property

rcomer · rcomer · commit 7345058c9887 · 2024-05-08T20:45:27.000+01:00
diff --git a/ci/mypy-stubtest-allowlist.txt b/ci/mypy-stubtest-allowlist.txt
@@ -46,3 +46,6 @@ matplotlib.tri.*TriInterpolator.gradient
 matplotlib.backend_bases.FigureCanvasBase._T
 matplotlib.backend_managers.ToolManager._T
 matplotlib.spines.Spine._T
+
+# Parameter inconsistency due to 3.10 deprecation
+matplotlib.figure.SubFigure.get_figure
diff --git a/doc/api/next_api_changes/behavior/28177-REC.rst b/doc/api/next_api_changes/behavior/28177-REC.rst
@@ -0,0 +1,7 @@
+SubFigure.get_figure
+~~~~~~~~~~~~~~~~~~~~
+
+...in future will by default return the direct parent figure, which may be a SubFigure.
+This will make the default behavior consistent with the
+`~matplotlib.artist.Artist.get_figure` method of other artists.  To control the
+behavior, use the *root* parameter.
diff --git a/lib/matplotlib/artist.py b/lib/matplotlib/artist.py
@@ -181,7 +181,7 @@ def __init__(self):
         self._stale = True
         self.stale_callback = None
         self._axes = None
-        self.figure = None
+        self._parent_figure = None
 
         self._transform = None
         self._transformSet = False
@@ -251,7 +251,7 @@ def remove(self):
             if self.figure:
                 if not _ax_flag:
                     self.figure.stale = True
-                self.figure = None
+                self._parent_figure = None
 
         else:
             raise NotImplementedError('cannot remove artist')
@@ -720,34 +720,49 @@ def set_path_effects(self, path_effects):
     def get_path_effects(self):
         return self._path_effects
 
-    def get_figure(self):
-        """Return the `.Figure` instance the artist belongs to."""
-        return self.figure
+    def get_figure(self, root=False):
+        """
+        Return the `.Figure` or `.SubFigure` instance the artist belongs to.
+
+        Parameters
+        ----------
+        root : bool, default=False
+            If False, return the (Sub)Figure this artist is on.  If True,
+            return the root Figure for a nested tree of SubFigures.
+        """
+        if root and self._parent_figure is not None:
+            return self._parent_figure.get_figure(root=True)
+
+        return self._parent_figure
 
     def set_figure(self, fig):
         """
-        Set the `.Figure` instance the artist belongs to.
+        Set the `.Figure` or `.SubFigure` instance the artist belongs to.
 
         Parameters
         ----------
-        fig : `~matplotlib.figure.Figure`
+        fig : `~matplotlib.figure.Figure` or `~matplotlib.figure.SubFigure`
         """
         # if this is a no-op just return
-        if self.figure is fig:
+        if self._parent_figure is fig:
             return
         # if we currently have a figure (the case of both `self.figure`
         # and *fig* being none is taken care of above) we then user is
         # trying to change the figure an artist is associated with which
         # is not allowed for the same reason as adding the same instance
         # to more than one Axes
-        if self.figure is not None:
+        if self._parent_figure is not None:
             raise RuntimeError("Can not put single artist in "
                                "more than one figure")
-        self.figure = fig
-        if self.figure and self.figure is not self:
+        self._parent_figure = fig
+        if self._parent_figure and self._parent_figure is not self:
             self.pchanged()
         self.stale = True
 
+    figure = property(get_figure, set_figure,
+                      doc=("The (Sub)Figure that the artist is on.  For more "
+                           "control, use the `get_figure` method."))
+
     def set_clip_box(self, clipbox):
         """
         Set the artist's clip `.Bbox`.
diff --git a/lib/matplotlib/artist.pyi b/lib/matplotlib/artist.pyi
@@ -87,8 +87,8 @@ class Artist:
     ) -> None: ...
     def set_path_effects(self, path_effects: list[AbstractPathEffect]) -> None: ...
     def get_path_effects(self) -> list[AbstractPathEffect]: ...
-    def get_figure(self) -> Figure | None: ...
-    def set_figure(self, fig: Figure) -> None: ...
+    def get_figure(self, root: bool = ...) -> Figure | SubFigure | None: ...
+    def set_figure(self, fig: Figure | SubFigure) -> None: ...
     def set_clip_box(self, clipbox: BboxBase | None) -> None: ...
     def set_clip_path(
         self,
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -1295,7 +1295,7 @@ def __clear(self):
         self._gridOn = mpl.rcParams['axes.grid']
         old_children, self._children = self._children, []
         for chld in old_children:
-            chld.axes = chld.figure = None
+            chld.axes = chld._parent_figure = None
         self._mouseover_set = _OrderedSet()
         self.child_axes = []
         self._current_image = None  # strictly for pyplot via _sci, _gci
diff --git a/lib/matplotlib/axes/_base.pyi b/lib/matplotlib/axes/_base.pyi
@@ -13,7 +13,7 @@ from matplotlib.cm import ScalarMappable
 from matplotlib.legend import Legend
 from matplotlib.lines import Line2D
 from matplotlib.gridspec import SubplotSpec, GridSpec
-from matplotlib.figure import Figure
+from matplotlib.figure import Figure, SubFigure
 from matplotlib.image import AxesImage
 from matplotlib.patches import Patch
 from matplotlib.scale import ScaleBase
@@ -81,7 +81,7 @@ class _AxesBase(martist.Artist):
     def get_subplotspec(self) -> SubplotSpec | None: ...
     def set_subplotspec(self, subplotspec: SubplotSpec) -> None: ...
     def get_gridspec(self) -> GridSpec | None: ...
-    def set_figure(self, fig: Figure) -> None: ...
+    def set_figure(self, fig: Figure | SubFigure) -> None: ...
     @property
     def viewLim(self) -> Bbox: ...
     def get_xaxis_transform(
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -30,6 +30,7 @@
 from contextlib import ExitStack
 import inspect
 import itertools
+import functools
 import logging
 from numbers import Integral
 import threading
@@ -227,6 +228,42 @@ def get_children(self):
                 *self.legends,
                 *self.subfigs]
 
+    def get_figure(self, root=None):
+        """
+        Return the `.Figure` or `.SubFigure` instance the (Sub)Figure belongs to.
+
+        Parameters
+        ----------
+        root : bool, default=True
+            If False, return the (Sub)Figure this artist is on.  If True,
+            return the root Figure for a nested tree of SubFigures.
+
+            .. deprecated:: 3.10
+
+                From version 3.12 *root* will default to False.
+        """
+        if self._root_figure is self:
+            # Top level Figure
+            return self
+
+        if root is None:
+            # When deprecation expires, consider removing the docstring and just
+            # inheriting the one from Artist.
+            message = ('From Matplotlib 3.12 SubFigure.get_figure will by default '
+                       'return the direct parent figure, which may be a SubFigure. '
+                       'To suppress this warning, pass the root parameter.')
+            _api.warn_deprecated('3.10', message=message)
+            root = True
+
+        if root:
+            return self._root_figure
+
+        return self._parent
+
+    figure = property(functools.partial(get_figure, root=True),
+                      doc=("The root `Figure`.  To get the parent of a `SubFigure`, "
+                           "use the `get_figure` method."))
+
     def contains(self, mouseevent):
         """
         Test whether the mouse event occurred on the figure.
@@ -2215,7 +2252,7 @@ def __init__(self, parent, subplotspec, *,
 
         self._subplotspec = subplotspec
         self._parent = parent
-        self.figure = parent.figure
+        self._root_figure = parent._root_figure
 
         # subfigures use the parent axstack
         self._axstack = parent._axstack
@@ -2493,7 +2530,7 @@ def __init__(self,
             %(Figure:kwdoc)s
         """
         super().__init__(**kwargs)
-        self.figure = self
+        self._root_figure = self
         self._layout_engine = None
 
         if layout is not None:
diff --git a/lib/matplotlib/offsetbox.pyi b/lib/matplotlib/offsetbox.pyi
@@ -2,7 +2,7 @@ import matplotlib.artist as martist
 from matplotlib.backend_bases import RendererBase, Event, FigureCanvasBase
 from matplotlib.colors import Colormap, Normalize
 import matplotlib.text as mtext
-from matplotlib.figure import Figure
+from matplotlib.figure import Figure, SubFigure
 from matplotlib.font_manager import FontProperties
 from matplotlib.image import BboxImage
 from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
@@ -26,7 +26,7 @@ class OffsetBox(martist.Artist):
     width: float | None
     height: float | None
     def __init__(self, *args, **kwargs) -> None: ...
-    def set_figure(self, fig: Figure) -> None: ...
+    def set_figure(self, fig: Figure | SubFigure) -> None: ...
     def set_offset(
         self,
         xy: tuple[float, float]
diff --git a/lib/matplotlib/quiver.pyi b/lib/matplotlib/quiver.pyi
@@ -1,7 +1,7 @@
 import matplotlib.artist as martist
 import matplotlib.collections as mcollections
 from matplotlib.axes import Axes
-from matplotlib.figure import Figure
+from matplotlib.figure import Figure, SubFigure
 from matplotlib.text import Text
 from matplotlib.transforms import Transform, Bbox
 
@@ -49,7 +49,7 @@ class QuiverKey(martist.Artist):
     ) -> None: ...
     @property
     def labelsep(self) -> float: ...
-    def set_figure(self, fig: Figure) -> None: ...
+    def set_figure(self, fig: Figure | SubFigure) -> None: ...
 
 class Quiver(mcollections.PolyCollection):
     X: ArrayLike
diff --git a/lib/matplotlib/tests/test_artist.py b/lib/matplotlib/tests/test_artist.py
@@ -562,3 +562,35 @@ def draw(self, renderer, extra):
 
     assert 'aardvark' == art.draw(renderer, 'aardvark')
     assert 'aardvark' == art.draw(renderer, extra='aardvark')
+
+
+def test_get_figure():
+    fig = plt.figure()
+    sfig1 = fig.subfigures()
+    sfig2 = sfig1.subfigures()
+    ax = sfig2.subplots()
+
+    assert fig.get_figure(root=True) is fig
+    assert fig.get_figure(root=False) is fig
+
+    assert ax.get_figure() is sfig2
+    assert ax.get_figure(root=False) is sfig2
+    assert ax.get_figure(root=True) is fig
+
+    # SubFigure.get_figure has separate implementation but should give consistent
+    # results to other artists.
+    assert sfig2.get_figure(root=False) is sfig1
+    assert sfig2.get_figure(root=True) is fig
+    # Currently different results by default.
+    with pytest.warns(mpl.MatplotlibDeprecationWarning):
+        assert sfig2.get_figure() is fig
+
+    # An artist not yet attached to anything has no figure.
+    ln = mlines.Line2D([], [])
+    assert ln.get_figure(root=True) is None
+    assert ln.get_figure(root=False) is None
+
+    # figure attribute is root for (Sub)Figures but parent for other artists.
+    assert ax.figure == sfig2
+    assert fig.figure == fig
+    assert sfig2.figure == fig
