Missing return types for Figure

ksunden · ksunden · commit 8edf04f929bf · 2023-04-27T12:38:12.000-05:00
Required changes for Figure.figure type hint
diff --git a/lib/matplotlib/artist.pyi b/lib/matplotlib/artist.pyi
@@ -4,7 +4,7 @@ from .axes._base import _AxesBase
 from .backend_bases import RendererBase, MouseEvent
 from .cm import ScalarMappable
 from .colors import BoundaryNorm
-from .figure import Figure, FigureBase
+from .figure import Figure, SubFigure
 from .path import Path
 from .patches import Patch
 from .patheffects import AbstractPathEffect
@@ -35,7 +35,7 @@ class Artist:
     zorder: float
     def __init_subclass__(cls): ...
     stale_callback: Callable[[Artist, bool], None] | None
-    figure: FigureBase | None
+    figure: Figure | SubFigure | None
     clipbox: Bbox | None
     def __init__(self) -> None: ...
     def remove(self) -> None: ...
diff --git a/lib/matplotlib/backend_bases.pyi b/lib/matplotlib/backend_bases.pyi
@@ -36,7 +36,7 @@ def get_registered_canvas_class(format: str) -> type[FigureCanvasBase]: ...
 
 class RendererBase:
     def __init__(self) -> None: ...
-    def open_group(self, s: str, gid: int | None = ...) -> None: ...
+    def open_group(self, s: str, gid: str | None = ...) -> None: ...
     def close_group(self, s: str) -> None: ...
     def draw_path(
         self,
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -188,7 +188,6 @@ def __init__(self, **kwargs):
         # axis._get_tick_boxes_siblings
         self._align_label_groups = {"x": cbook.Grouper(), "y": cbook.Grouper()}
 
-        self.figure = self
         self._localaxes = []  # track all axes
         self.artists = []
         self.lines = []
@@ -2458,6 +2457,7 @@ def __init__(self,
             %(Figure:kwdoc)s
         """
         super().__init__(**kwargs)
+        self.figure = self
         self._layout_engine = None
 
         if layout is not None:
diff --git a/lib/matplotlib/figure.pyi b/lib/matplotlib/figure.pyi
@@ -15,7 +15,7 @@ from matplotlib.colors import Colormap, Normalize
 from matplotlib.colorbar import Colorbar
 from matplotlib.cm import ScalarMappable
 from matplotlib.gridspec import GridSpec, SubplotSpec
-from matplotlib.image import _ImageBase
+from matplotlib.image import _ImageBase, FigureImage
 from matplotlib.layout_engine import (
     ConstrainedLayoutEngine,
     LayoutEngine,
@@ -69,7 +69,6 @@ class SubplotParams:
     ) -> None: ...
 
 class FigureBase(Artist):
-    figure: FigureBase | None
     artists: list[Artist]
     lines: list[Line2D]
     patches: list[Patch]
@@ -154,15 +153,15 @@ class FigureBase(Artist):
         *,
         sharex: bool | Literal["none", "all", "row", "col"] = ...,
         sharey: bool | Literal["none", "all", "row", "col"] = ...,
-        squeeze: Literal[True] = ...,
+        squeeze: bool = ...,
         width_ratios: ArrayLike | None = ...,
         height_ratios: ArrayLike | None = ...,
         subplot_kw: dict[str, Any] | None = ...,
         gridspec_kw: dict[str, Any] | None = ...
     ) -> np.ndarray | SubplotBase | Axes: ...
     def delaxes(self, ax: Axes) -> None: ...
     def clear(self, keep_observers: bool = ...) -> None: ...
-    def clf(self, keep_observers: bool = ...): ...
+    def clf(self, keep_observers: bool = ...) -> None: ...
 
     @overload
     def legend(self) -> Legend: ...
@@ -203,7 +202,7 @@ class FigureBase(Artist):
     def align_xlabels(self, axs: Iterable[Axes] | None = ...) -> None: ...
     def align_ylabels(self, axs: Iterable[Axes] | None = ...) -> None: ...
     def align_labels(self, axs: Iterable[Axes] | None = ...) -> None: ...
-    def add_gridspec(self, nrows: int = ..., ncols: int = ..., **kwargs): ...
+    def add_gridspec(self, nrows: int = ..., ncols: int = ..., **kwargs) -> GridSpec: ...
     @overload
     def subfigures(
         self,
@@ -234,7 +233,7 @@ class FigureBase(Artist):
     def _gci(self) -> ScalarMappable | None: ...
     def _process_projection_requirements(
         self, *args, axes_class=None, polar=False, projection=None, **kwargs
-    ): ...
+    ) -> tuple[type[Axes], dict[str, Any]]: ...
     def get_default_bbox_extra_artists(self) -> list[Artist]: ...
     def get_tightbbox(
         self,
@@ -258,7 +257,7 @@ class FigureBase(Artist):
     ) -> dict[Any, Axes]: ...
 
 class SubFigure(FigureBase):
-    figure: FigureBase
+    figure: Figure
     subplotpars: SubplotParams
     dpi_scale_trans: Affine2D
     canvas: FigureCanvasBase
@@ -295,6 +294,7 @@ class SubFigure(FigureBase):
     def get_axes(self) -> list[Axes]: ...
 
 class Figure(FigureBase):
+    figure: Figure
     bbox_inches: Bbox
     dpi_scale_trans: Affine2D
     bbox: Bbox
@@ -357,7 +357,7 @@ class Figure(FigureBase):
         origin: Literal["upper", "lower"] | None = ...,
         resize: bool = ...,
         **kwargs
-    ): ...
+    ) -> FigureImage: ...
     def set_size_inches(
         self, w: float | tuple[float, float], h: float | None = ..., forward: bool = ...
     ) -> None: ...
@@ -371,7 +371,7 @@ class Figure(FigureBase):
     def clear(self, keep_observers: bool = ...) -> None: ...
     def draw_without_rendering(self) -> None: ...
     def draw_artist(self, a: Artist) -> None: ...
-    def add_axobserver(self, func: Callable[[Figure], Any]): ...
+    def add_axobserver(self, func: Callable[[Figure], Any]) -> None: ...
     def savefig(
         self,
         fname: str | os.PathLike | IO,
@@ -388,7 +388,7 @@ class Figure(FigureBase):
         mouse_pop: MouseButton = ...,
         mouse_stop: MouseButton = ...,
     ) -> list[tuple[int, int]]: ...
-    def waitforbuttonpress(self, timeout: float = ...): ...
+    def waitforbuttonpress(self, timeout: float = ...) -> None | bool: ...
     def tight_layout(
         self,
         *,
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -116,7 +116,7 @@
     from matplotlib.figure import SubFigure
     from matplotlib.legend import Legend
     from matplotlib.mlab import GaussianKDE
-    from matplotlib.image import AxesImage
+    from matplotlib.image import AxesImage, FigureImage
     from matplotlib.patches import FancyArrow, StepPatch
     from matplotlib.quiver import Barbs, Quiver, QuiverKey
     from matplotlib.scale import ScaleBase
@@ -754,7 +754,7 @@ def figure(
     FigureClass: Type[Figure] = Figure,
     clear: bool = False,
     **kwargs
-):
+) -> Figure:
     """
     Create a new figure, or activate an existing figure.
 
@@ -857,6 +857,7 @@ def figure(
     in the matplotlibrc file.
     """
     if isinstance(num, FigureBase):
+        # type narrowed to `Figure | SubFigure` by combination of input and isinstance
         if num.canvas.manager is None:
             raise ValueError("The passed figure is not managed by pyplot")
         _pylab_helpers.Gcf.set_active(num.canvas.manager)
@@ -2347,7 +2348,7 @@ def matshow(A: ArrayLike, fignum: None | int = None, **kwargs) -> AxesImage:
         # Extract actual aspect ratio of array and make appropriately sized
         # figure.
         fig = figure(fignum, figsize=figaspect(A))
-        ax = fig.add_axes([0.15, 0.09, 0.775, 0.775])
+        ax = fig.add_axes((0.15, 0.09, 0.775, 0.775))
     im = ax.matshow(A, **kwargs)
     sci(im)
     return im
@@ -2403,7 +2404,7 @@ def figimage(
     origin: Literal["upper", "lower"] | None = None,
     resize: bool = False,
     **kwargs,
-):
+) -> FigureImage:
     return gcf().figimage(
         X,
         xo=xo,
@@ -2494,7 +2495,7 @@ def tight_layout(
 
 # Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
 @_copy_docstring_and_deprecators(Figure.waitforbuttonpress)
-def waitforbuttonpress(timeout: float = -1):
+def waitforbuttonpress(timeout: float = -1) -> None | bool:
     return gcf().waitforbuttonpress(timeout=timeout)
 
 
