Complete seaborn stubs and mark as strict (#10974)

stubs/seaborn/METADATA.toml

@@ -4,6 +4,3 @@ requires = ["matplotlib>=3.8", "numpy>=1.20", "pandas-stubs"]
 # matplotlib>=3.8 requires Python >=3.9
 requires_python = ">=3.9"
 upstream_repository = "https://github.com/mwaskom/seaborn"
-
-[tool.stubtest]
-ignore_missing_stub = false

stubs/seaborn/seaborn/_core/groupby.pyi

@@ -1,10 +1,33 @@
 from _typeshed import Incomplete
-from collections.abc import Callable
+from collections.abc import Callable, Hashable, Mapping
+from typing_extensions import Concatenate, ParamSpec, TypeAlias
 
+from numpy import ufunc
 from pandas import DataFrame
 
+# pandas._typing.AggFuncTypeFrame is partially Unknown
+_AggFuncTypeBase: TypeAlias = Callable[..., Incomplete] | str | ufunc
+# Using Hashable instead of HashableT to work around pytype issue
+_AggFuncTypeDictFrame: TypeAlias = Mapping[Hashable, _AggFuncTypeBase | list[_AggFuncTypeBase]]
+_AggFuncTypeFrame: TypeAlias = _AggFuncTypeBase | list[_AggFuncTypeBase] | _AggFuncTypeDictFrame
+
+_P = ParamSpec("_P")
+
 class GroupBy:
     order: dict[str, list[Incomplete] | None]
     def __init__(self, order: list[str] | dict[str, list[Incomplete] | None]) -> None: ...
-    def agg(self, data: DataFrame, *args, **kwargs) -> DataFrame: ...
-    def apply(self, data: DataFrame, func: Callable[..., DataFrame], *args, **kwargs) -> DataFrame: ...
+    # Signature based on pandas.core.groupby.generic.DataFrameGroupBy.aggregate
+    # args and kwargs possible values depend on func which itself can be
+    # an attribute name, a mapping, a callable, or lead to a jitted numba function
+    def agg(
+        self,
+        data: DataFrame,
+        func: _AggFuncTypeFrame = ...,
+        *args: Incomplete,
+        engine: str | None = None,
+        engine_kwargs: dict[str, bool] | None = None,
+        **kwargs: Incomplete,
+    ) -> DataFrame: ...
+    def apply(
+        self, data: DataFrame, func: Callable[Concatenate[DataFrame, _P], DataFrame], *args: _P.args, **kwargs: _P.kwargs
+    ) -> DataFrame: ...

stubs/seaborn/seaborn/_core/plot.pyi

@@ -4,11 +4,14 @@ from _typeshed import Incomplete, SupportsKeysAndGetItem
 from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from typing import IO, Any, NoReturn, TypeVar
-from typing_extensions import Literal, TypedDict
+from typing_extensions import Literal, Never, Self, TypedDict
 
 import matplotlib as mpl
+from matplotlib.artist import Artist
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure, SubFigure
+from matplotlib.transforms import Bbox
+from matplotlib.typing import ColorType
 from seaborn._core.data import PlotData
 from seaborn._core.moves import Move
 from seaborn._core.scales import Scale
@@ -69,7 +72,7 @@ class Plot:
     __signature__: inspect.Signature
     config: PlotConfig
     def __init__(self, *args: DataSource | VariableSpec, data: DataSource = None, **variables: VariableSpec) -> None: ...
-    def __add__(self, other) -> NoReturn: ...
+    def __add__(self, other: Never) -> NoReturn: ...
     def on(self, target: Axes | SubFigure | Figure) -> Plot: ...
     def add(
         self,
@@ -97,11 +100,48 @@ class Plot:
     def label(self, *, title: str | None = None, legend: str | None = None, **variables: str | Callable[[str], str]) -> Plot: ...
     def layout(self, *, size: tuple[float, float] | Default = ..., engine: str | None | Default = ...) -> Plot: ...
     def theme(self, *args: dict[str, Any]) -> Plot: ...
-    def save(self, loc: str | os.PathLike[Any] | IO[Any], **kwargs) -> Plot: ...
-    def show(self, **kwargs) -> None: ...
+    # Same signature as Plotter.save
+    def save(
+        self,
+        loc: str | os.PathLike[Any] | IO[Any],
+        *,
+        transparent: bool | None = None,
+        dpi: float | None = 96,
+        facecolor: ColorType | Literal["auto"] | None = ...,
+        edgecolor: ColorType | Literal["auto"] | None = ...,
+        orientation: str = "protrait",
+        format: str | None = None,
+        bbox_inches: Literal["tight"] | Bbox | None = None,
+        pad_inches: float | None = None,
+        bbox_extra_artists: list[Artist] | None = None,
+        backend: str | None = None,
+        **kwargs: Any,
+    ) -> Self: ...
+    # Same signature as Plotter.show
+    def show(self, *, block: bool | None = None) -> None: ...
     def plot(self, pyplot: bool = False) -> Plotter: ...
 
 class Plotter:
     def __init__(self, pyplot: bool, theme: dict[str, Any]) -> None: ...
-    def save(self, loc: str | os.PathLike[Any] | IO[Any], **kwargs) -> Plotter: ...
-    def show(self, **kwargs) -> None: ...
+    def save(
+        self,
+        loc: str | os.PathLike[Any] | IO[Any],
+        *,
+        # From matplotlib.figure.Figure.savefig
+        transparent: bool | None = None,
+        # keyword-only arguments below are the same as matplotlib.backend_bases.FigureCanvasBase.print_figure
+        # but with different defaults
+        dpi: float | None = 96,
+        facecolor: ColorType | Literal["auto"] | None = ...,
+        edgecolor: ColorType | Literal["auto"] | None = ...,
+        orientation: str = "protrait",
+        format: str | None = None,
+        bbox_inches: Literal["tight"] | Bbox | None = None,
+        pad_inches: float | None = None,
+        bbox_extra_artists: list[Artist] | None = None,
+        backend: str | None = None,
+        # Further **kwargs can truly be anything from an overriden Canvas method that is still passed down
+        **kwargs: Any,
+    ) -> Self: ...
+    # Same as matplotlib.backend_bases._Bases. No other backend override show
+    def show(self, *, block: bool | None = None) -> None: ...

stubs/seaborn/seaborn/external/docscrape.pyi

@@ -44,7 +44,7 @@ def dedent_lines(lines: Iterable[str]) -> list[str]: ...
 def header(text: str, style: str = "-") -> str: ...
 
 class FunctionDoc(NumpyDocString):
-    def __init__(self, func, role: str = "func", doc: str | None = None, config: Unused = {}) -> None: ...
+    def __init__(self, func: object, role: str = "func", doc: str | None = None, config: Unused = {}) -> None: ...
     def get_func(self) -> tuple[Incomplete, str]: ...
 
 class ClassDoc(NumpyDocString):

stubs/seaborn/seaborn/matrix.pyi

@@ -1,21 +1,25 @@
 from _typeshed import Incomplete
-from collections.abc import Mapping, Sequence
-from typing_extensions import Literal, Self
+from collections.abc import Hashable, Iterable, Mapping, Sequence
+from typing_extensions import Literal, Self, TypeAlias
 
 import numpy as np
 from matplotlib.axes import Axes
 from matplotlib.colors import Colormap, ListedColormap
 from matplotlib.gridspec import GridSpec
 from matplotlib.typing import ColorType
+from numpy._typing import _ArrayLike, _ArrayLikeInt_co
 from numpy.typing import ArrayLike, NDArray
-from pandas import DataFrame
+from pandas import DataFrame, Index, Series
 
 from .axisgrid import Grid
 
+# pandas._typing.ListLikeU is partially Unknown
+_ListLikeU: TypeAlias = Sequence[Incomplete] | np.ndarray[Incomplete, Incomplete] | Series[Incomplete] | Index[Incomplete]
+
 __all__ = ["heatmap", "clustermap"]
 
 def heatmap(
-    data: Incomplete,
+    data: DataFrame | _ArrayLike[Incomplete],
     *,
     vmin: float | None = None,
     vmax: float | None = None,
@@ -101,7 +105,11 @@ class ClusterGrid(Grid):
     dendrogram_col: _DendrogramPlotter | None
     def __init__(
         self,
-        data: Incomplete,
+        data: _ListLikeU
+        | DataFrame
+        | dict[Incomplete, Incomplete]
+        | Iterable[_ListLikeU | tuple[Hashable, _ListLikeU] | dict[Incomplete, Incomplete]]
+        | None,
         pivot_kws: Mapping[str, Incomplete] | None = None,
         z_score: int | None = None,
         standard_scale: int | None = None,
@@ -124,10 +132,10 @@ class ClusterGrid(Grid):
     def z_score(data2d: DataFrame, axis: int = 1) -> DataFrame: ...
     @staticmethod
     def standard_scale(data2d: DataFrame, axis: int = 1) -> DataFrame: ...
-    def dim_ratios(self, colors: Incomplete, dendrogram_ratio: float, colors_ratio: float) -> list[float]: ...
+    def dim_ratios(self, colors: Incomplete | None, dendrogram_ratio: float, colors_ratio: float) -> list[float]: ...
     @staticmethod
     def color_list_to_matrix_and_cmap(
-        colors: Sequence[ColorType], ind: list[int], axis: int = 0
+        colors: Sequence[ColorType], ind: _ArrayLikeInt_co, axis: int = 0
     ) -> tuple[NDArray[np.int_], ListedColormap]: ...
     def plot_dendrograms(
         self,
@@ -139,8 +147,10 @@ class ClusterGrid(Grid):
         col_linkage: NDArray[Incomplete] | None,
         tree_kws: dict[str, Incomplete] | None,
     ) -> None: ...
-    def plot_colors(self, xind: Incomplete, yind: Incomplete, **kws: Incomplete) -> None: ...
-    def plot_matrix(self, colorbar_kws: dict[str, Incomplete], xind: Incomplete, yind: Incomplete, **kws: Incomplete) -> None: ...
+    def plot_colors(self, xind: _ArrayLikeInt_co, yind: _ArrayLikeInt_co, **kws: Incomplete) -> None: ...
+    def plot_matrix(
+        self, colorbar_kws: dict[str, Incomplete], xind: _ArrayLikeInt_co, yind: _ArrayLikeInt_co, **kws: Incomplete
+    ) -> None: ...
     def plot(
         self,
         metric: str,
@@ -155,7 +165,11 @@ class ClusterGrid(Grid):
     ) -> Self: ...
 
 def clustermap(
-    data: Incomplete,
+    data: _ListLikeU
+    | DataFrame
+    | dict[Incomplete, Incomplete]
+    | Iterable[_ListLikeU | tuple[Hashable, _ListLikeU] | dict[Incomplete, Incomplete]]
+    | None,
     *,
     pivot_kws: dict[str, Incomplete] | None = None,
     method: str = "average",

stubs/seaborn/seaborn/miscplot.pyi

@@ -1,3 +1,4 @@
+from _typeshed import Unused
 from collections.abc import Sequence
 
 from matplotlib.typing import ColorType
@@ -5,4 +6,4 @@ from matplotlib.typing import ColorType
 __all__ = ["palplot", "dogplot"]
 
 def palplot(pal: Sequence[ColorType], size: int = 1) -> None: ...
-def dogplot(*_, **__) -> None: ...
+def dogplot(*_: Unused, **__: Unused) -> None: ...

stubs/seaborn/seaborn/rcmod.pyi

@@ -1,3 +1,4 @@
+from _typeshed import Unused
 from collections.abc import Callable, Sequence
 from typing import Any, TypeVar
 from typing_extensions import Literal
@@ -56,7 +57,7 @@ def set_context(
 
 class _RCAesthetics(dict[_KT, _VT]):
     def __enter__(self) -> None: ...
-    def __exit__(self, exc_type, exc_value, exc_tb) -> None: ...
+    def __exit__(self, exc_type: Unused, exc_value: Unused, exc_tb: Unused) -> None: ...
     def __call__(self, func: _F) -> _F: ...
 
 class _AxesStyle(_RCAesthetics[_KT, _VT]): ...