From 613a8e88d6753fcc985760395db35b5d951aa41d Mon Sep 17 00:00:00 2001 From: "Adam J. Stewart" Date: Sat, 6 Jul 2024 17:51:53 +0200 Subject: [PATCH 1/5] [TYP] Fix overload of `pyplot.subplots` --- lib/matplotlib/figure.pyi | 14 ++++++++++++++ lib/matplotlib/pyplot.py | 17 +++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/lib/matplotlib/figure.pyi b/lib/matplotlib/figure.pyi index b079312695c1..91196f6add8e 100644 --- a/lib/matplotlib/figure.pyi +++ b/lib/matplotlib/figure.pyi @@ -106,6 +106,20 @@ class FigureBase(Artist): gridspec_kw: dict[str, Any] | None = ..., ) -> Axes: ... @overload + def subplots( + self, + nrows: int = ..., + ncols: int = ..., + *, + sharex: bool | Literal["none", "all", "row", "col"] = ..., + sharey: bool | Literal["none", "all", "row", "col"] = ..., + squeeze: Literal[True], + width_ratios: Sequence[float] | None = ..., + height_ratios: Sequence[float] | None = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + ) -> np.ndarray: ... # TODO numpy/numpy#24738 + @overload def subplots( self, nrows: int = ..., diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py index 441af598dbc6..f196a8e9d440 100644 --- a/lib/matplotlib/pyplot.py +++ b/lib/matplotlib/pyplot.py @@ -1579,6 +1579,23 @@ def subplots( ... +@overload +def subplots( + nrows: int = ..., + ncols: int = ..., + *, + sharex: bool | Literal["none", "all", "row", "col"] = ..., + sharey: bool | Literal["none", "all", "row", "col"] = ..., + squeeze: Literal[True] = ..., + width_ratios: Sequence[float] | None = ..., + height_ratios: Sequence[float] | None = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + **fig_kw +) -> tuple[Figure, np.ndarray]: # TODO numpy/numpy#24738 + ... + + @overload def subplots( nrows: int = ..., From 23adb361b96f628293f9d56464ec312921536b93 Mon Sep 17 00:00:00 2001 From: "Adam J. Stewart" Date: Wed, 10 Jul 2024 19:54:37 +0200 Subject: [PATCH 2/5] Simplify the stub --- lib/matplotlib/figure.pyi | 30 +----------------------------- 1 file changed, 1 insertion(+), 29 deletions(-) diff --git a/lib/matplotlib/figure.pyi b/lib/matplotlib/figure.pyi index 91196f6add8e..13e2adfe1b69 100644 --- a/lib/matplotlib/figure.pyi +++ b/lib/matplotlib/figure.pyi @@ -92,34 +92,6 @@ class FigureBase(Artist): @overload def add_subplot(self, **kwargs) -> Axes: ... @overload - def subplots( - self, - nrows: Literal[1] = ..., - ncols: Literal[1] = ..., - *, - sharex: bool | Literal["none", "all", "row", "col"] = ..., - sharey: bool | Literal["none", "all", "row", "col"] = ..., - squeeze: Literal[True] = ..., - width_ratios: Sequence[float] | None = ..., - height_ratios: Sequence[float] | None = ..., - subplot_kw: dict[str, Any] | None = ..., - gridspec_kw: dict[str, Any] | None = ..., - ) -> Axes: ... - @overload - def subplots( - self, - nrows: int = ..., - ncols: int = ..., - *, - sharex: bool | Literal["none", "all", "row", "col"] = ..., - sharey: bool | Literal["none", "all", "row", "col"] = ..., - squeeze: Literal[True], - width_ratios: Sequence[float] | None = ..., - height_ratios: Sequence[float] | None = ..., - subplot_kw: dict[str, Any] | None = ..., - gridspec_kw: dict[str, Any] | None = ..., - ) -> np.ndarray: ... # TODO numpy/numpy#24738 - @overload def subplots( self, nrows: int = ..., @@ -146,7 +118,7 @@ class FigureBase(Artist): height_ratios: Sequence[float] | None = ..., subplot_kw: dict[str, Any] | None = ..., gridspec_kw: dict[str, Any] | None = ..., - ) -> Axes | np.ndarray: ... + ) -> Any: ... def delaxes(self, ax: Axes) -> None: ... def clear(self, keep_observers: bool = ...) -> None: ... def clf(self, keep_observers: bool = ...) -> None: ... From 64cac364f678c5d34e1b12da6a02c231b8eb2f7a Mon Sep 17 00:00:00 2001 From: "Adam J. Stewart" Date: Wed, 10 Jul 2024 21:58:50 +0200 Subject: [PATCH 3/5] Simplify the code --- lib/matplotlib/pyplot.py | 51 ---------------------------------------- 1 file changed, 51 deletions(-) diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py index f196a8e9d440..b0ffe23efeea 100644 --- a/lib/matplotlib/pyplot.py +++ b/lib/matplotlib/pyplot.py @@ -1562,40 +1562,6 @@ def subplot(*args, **kwargs) -> Axes: return ax -@overload -def subplots( - nrows: Literal[1] = ..., - ncols: Literal[1] = ..., - *, - sharex: bool | Literal["none", "all", "row", "col"] = ..., - sharey: bool | Literal["none", "all", "row", "col"] = ..., - squeeze: Literal[True] = ..., - width_ratios: Sequence[float] | None = ..., - height_ratios: Sequence[float] | None = ..., - subplot_kw: dict[str, Any] | None = ..., - gridspec_kw: dict[str, Any] | None = ..., - **fig_kw -) -> tuple[Figure, Axes]: - ... - - -@overload -def subplots( - nrows: int = ..., - ncols: int = ..., - *, - sharex: bool | Literal["none", "all", "row", "col"] = ..., - sharey: bool | Literal["none", "all", "row", "col"] = ..., - squeeze: Literal[True] = ..., - width_ratios: Sequence[float] | None = ..., - height_ratios: Sequence[float] | None = ..., - subplot_kw: dict[str, Any] | None = ..., - gridspec_kw: dict[str, Any] | None = ..., - **fig_kw -) -> tuple[Figure, np.ndarray]: # TODO numpy/numpy#24738 - ... - - @overload def subplots( nrows: int = ..., @@ -1613,23 +1579,6 @@ def subplots( ... -@overload -def subplots( - nrows: int = ..., - ncols: int = ..., - *, - sharex: bool | Literal["none", "all", "row", "col"] = ..., - sharey: bool | Literal["none", "all", "row", "col"] = ..., - squeeze: bool = ..., - width_ratios: Sequence[float] | None = ..., - height_ratios: Sequence[float] | None = ..., - subplot_kw: dict[str, Any] | None = ..., - gridspec_kw: dict[str, Any] | None = ..., - **fig_kw -) -> tuple[Figure, Axes | np.ndarray]: - ... - - def subplots( nrows: int = 1, ncols: int = 1, *, sharex: bool | Literal["none", "all", "row", "col"] = False, From c7a128eec1123866183d412947270ab2c2b11ebc Mon Sep 17 00:00:00 2001 From: "Adam J. Stewart" Date: Wed, 10 Jul 2024 22:06:48 +0200 Subject: [PATCH 4/5] Code needs two overloads --- lib/matplotlib/pyplot.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py index b0ffe23efeea..9660f17ddac1 100644 --- a/lib/matplotlib/pyplot.py +++ b/lib/matplotlib/pyplot.py @@ -1562,6 +1562,23 @@ def subplot(*args, **kwargs) -> Axes: return ax +@overload +def subplots( + nrows: int = ..., + ncols: int = ..., + *, + sharex: bool | Literal["none", "all", "row", "col"] = ..., + sharey: bool | Literal["none", "all", "row", "col"] = ..., + squeeze: Literal[True], + width_ratios: Sequence[float] | None = ..., + height_ratios: Sequence[float] | None = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + **fig_kw +) -> tuple[Figure, Any]: + ... + + @overload def subplots( nrows: int = ..., From 4242506cd51bb635af41d9fb10a08469bca66d46 Mon Sep 17 00:00:00 2001 From: "Adam J. Stewart" Date: Thu, 11 Jul 2024 10:15:14 +0200 Subject: [PATCH 5/5] Avoid union for dynamic type hints --- lib/matplotlib/figure.pyi | 14 ++++++++++++++ lib/matplotlib/pyplot.py | 25 +++++++++++++++++++++---- 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/lib/matplotlib/figure.pyi b/lib/matplotlib/figure.pyi index 13e2adfe1b69..c31f90b4b2a8 100644 --- a/lib/matplotlib/figure.pyi +++ b/lib/matplotlib/figure.pyi @@ -92,6 +92,20 @@ class FigureBase(Artist): @overload def add_subplot(self, **kwargs) -> Axes: ... @overload + def subplots( + self, + nrows: Literal[1] = ..., + ncols: Literal[1] = ..., + *, + sharex: bool | Literal["none", "all", "row", "col"] = ..., + sharey: bool | Literal["none", "all", "row", "col"] = ..., + squeeze: Literal[True] = ..., + width_ratios: Sequence[float] | None = ..., + height_ratios: Sequence[float] | None = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + ) -> Axes: ... + @overload def subplots( self, nrows: int = ..., diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py index 9660f17ddac1..9587850173d6 100644 --- a/lib/matplotlib/pyplot.py +++ b/lib/matplotlib/pyplot.py @@ -1564,18 +1564,18 @@ def subplot(*args, **kwargs) -> Axes: @overload def subplots( - nrows: int = ..., - ncols: int = ..., + nrows: Literal[1] = ..., + ncols: Literal[1] = ..., *, sharex: bool | Literal["none", "all", "row", "col"] = ..., sharey: bool | Literal["none", "all", "row", "col"] = ..., - squeeze: Literal[True], + squeeze: Literal[True] = ..., width_ratios: Sequence[float] | None = ..., height_ratios: Sequence[float] | None = ..., subplot_kw: dict[str, Any] | None = ..., gridspec_kw: dict[str, Any] | None = ..., **fig_kw -) -> tuple[Figure, Any]: +) -> tuple[Figure, Axes]: ... @@ -1596,6 +1596,23 @@ def subplots( ... +@overload +def subplots( + nrows: int = ..., + ncols: int = ..., + *, + sharex: bool | Literal["none", "all", "row", "col"] = ..., + sharey: bool | Literal["none", "all", "row", "col"] = ..., + squeeze: bool = ..., + width_ratios: Sequence[float] | None = ..., + height_ratios: Sequence[float] | None = ..., + subplot_kw: dict[str, Any] | None = ..., + gridspec_kw: dict[str, Any] | None = ..., + **fig_kw +) -> tuple[Figure, Any]: + ... + + def subplots( nrows: int = 1, ncols: int = 1, *, sharex: bool | Literal["none", "all", "row", "col"] = False,