From e109588be023d197bbc2c3311f00d80708b28d63 Mon Sep 17 00:00:00 2001 From: Avasam Date: Wed, 15 Nov 2023 06:58:27 -0500 Subject: [PATCH] Complete `seaborn` stubs and mark as strict (#10974) --- pyrightconfig.stricter.json | 1 - stubs/seaborn/METADATA.toml | 3 -- stubs/seaborn/seaborn/_core/groupby.pyi | 29 +++++++++-- stubs/seaborn/seaborn/_core/plot.pyi | 52 +++++++++++++++++--- stubs/seaborn/seaborn/external/docscrape.pyi | 2 +- stubs/seaborn/seaborn/matrix.pyi | 34 +++++++++---- stubs/seaborn/seaborn/miscplot.pyi | 3 +- stubs/seaborn/seaborn/rcmod.pyi | 3 +- 8 files changed, 101 insertions(+), 26 deletions(-) diff --git a/pyrightconfig.stricter.json b/pyrightconfig.stricter.json index 493814390..7975ccb0d 100644 --- a/pyrightconfig.stricter.json +++ b/pyrightconfig.stricter.json @@ -70,7 +70,6 @@ "stubs/redis", "stubs/requests", "stubs/requests-oauthlib", - "stubs/seaborn", "stubs/setuptools", "stubs/stripe", "stubs/tqdm", diff --git a/stubs/seaborn/METADATA.toml b/stubs/seaborn/METADATA.toml index b8032f4a4..1aca3a6bd 100644 --- a/stubs/seaborn/METADATA.toml +++ b/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 diff --git a/stubs/seaborn/seaborn/_core/groupby.pyi b/stubs/seaborn/seaborn/_core/groupby.pyi index bfbe4c469..41619174b 100644 --- a/stubs/seaborn/seaborn/_core/groupby.pyi +++ b/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: ... diff --git a/stubs/seaborn/seaborn/_core/plot.pyi b/stubs/seaborn/seaborn/_core/plot.pyi index 18981477d..eaed6ab60 100644 --- a/stubs/seaborn/seaborn/_core/plot.pyi +++ b/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: ... diff --git a/stubs/seaborn/seaborn/external/docscrape.pyi b/stubs/seaborn/seaborn/external/docscrape.pyi index f8f936e35..7c2c73f05 100644 --- a/stubs/seaborn/seaborn/external/docscrape.pyi +++ b/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): diff --git a/stubs/seaborn/seaborn/matrix.pyi b/stubs/seaborn/seaborn/matrix.pyi index 525096746..db5d8d4d0 100644 --- a/stubs/seaborn/seaborn/matrix.pyi +++ b/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", diff --git a/stubs/seaborn/seaborn/miscplot.pyi b/stubs/seaborn/seaborn/miscplot.pyi index 7ec684217..82e8dbac5 100644 --- a/stubs/seaborn/seaborn/miscplot.pyi +++ b/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: ... diff --git a/stubs/seaborn/seaborn/rcmod.pyi b/stubs/seaborn/seaborn/rcmod.pyi index 3cecf9497..f80c5ae8d 100644 --- a/stubs/seaborn/seaborn/rcmod.pyi +++ b/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]): ...