[TYP] Change overload of pyplot.subplots

n-takumasa · n-takumasa · commit 07ec432b9c98 · 2023-10-05T23:13:25.000Z
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -1440,48 +1440,26 @@ def subplot(*args, **kwargs) -> Axes:
     return ax
 
 
-# NOTE  The actual type is `Axes` or `numpy.ndarray`,
-# but `numpy.ndarray` does notsupport objects.
-# NOTE Since there is no Exclude-type in Python's type hints, it is assumed that
-# the overload that matches first will be resolved.
-# mypy warns that it is an unsafe overload, so mark it as ignore.
 @overload  # type: ignore[misc]
 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  # type: ignore[misc]
-def subplots(
-    nrows: Literal[1] = ...,
+    nrows: int = ...,
     ncols: int = ...,
     *,
     sharex: bool | Literal["none", "all", "row", "col"] = ...,
     sharey: bool | Literal["none", "all", "row", "col"] = ...,
-    squeeze: Literal[True] = ...,
+    squeeze: Literal[False] = ...,
     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, Sequence[Axes]]:
+) -> tuple[Figure, np.ndarray]:  # TODO numpy/numpy#24738
     ...
 
 
 @overload  # type: ignore[misc]
 def subplots(
-    nrows: int = ...,
+    nrows: Literal[1] = ...,
     ncols: Literal[1] = ...,
     *,
     sharex: bool | Literal["none", "all", "row", "col"] = ...,
@@ -1492,24 +1470,24 @@ def subplots(
     subplot_kw: dict[str, Any] | None = ...,
     gridspec_kw: dict[str, Any] | None = ...,
     **fig_kw
-) -> tuple[Figure, Sequence[Axes]]:
+) -> tuple[Figure, Axes]:
     ...
 
 
-@overload  # type: ignore[misc]
+@overload
 def subplots(
     nrows: int = ...,
     ncols: int = ...,
     *,
     sharex: bool | Literal["none", "all", "row", "col"] = ...,
     sharey: bool | Literal["none", "all", "row", "col"] = ...,
-    squeeze: Literal[False] = ...,
+    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, Sequence[Sequence[Axes]]]:
+) -> tuple[Figure, Axes | np.ndarray]:
     ...
 
 
