From 7864edbc397ec90c51198664c6487ed903e05802 Mon Sep 17 00:00:00 2001 From: Ali Hamdan Date: Sat, 26 Jul 2025 17:42:14 +0200 Subject: [PATCH] seaborn: improve the matrix module (#14464) --- stubs/seaborn/seaborn/matrix.pyi | 139 ++++++++++++++++++------------- 1 file changed, 81 insertions(+), 58 deletions(-) diff --git a/stubs/seaborn/seaborn/matrix.pyi b/stubs/seaborn/seaborn/matrix.pyi index 3a9a25e54..34ba5bb99 100644 --- a/stubs/seaborn/seaborn/matrix.pyi +++ b/stubs/seaborn/seaborn/matrix.pyi @@ -1,26 +1,41 @@ from _typeshed import Incomplete from collections.abc import Hashable, Iterable, Mapping, Sequence -from typing import Literal +from typing import Literal, TypedDict, type_check_only from typing_extensions import Self, TypeAlias import numpy as np +import pandas as pd from matplotlib.axes import Axes -from matplotlib.colors import Colormap, ListedColormap +from matplotlib.colors import Colormap, ListedColormap, Normalize from matplotlib.gridspec import GridSpec from matplotlib.typing import ColorType -from numpy._typing import _ArrayLike, _ArrayLikeInt_co +from numpy._typing import _ArrayLikeInt_co from numpy.typing import ArrayLike, NDArray -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] +_ListLikeU: TypeAlias = Sequence[Incomplete] | NDArray[Incomplete] | pd.Series[Incomplete] | pd.Index[Incomplete] +_ConvertibleToDataFrame: TypeAlias = ( + _ListLikeU + | pd.DataFrame + | dict[Incomplete, Incomplete] + | Iterable[_ListLikeU | tuple[Hashable, _ListLikeU] | dict[Incomplete, Incomplete]] + | None +) +_FlatOrNestedSequenceOfColors: TypeAlias = ( + Sequence[ColorType] + | Sequence[Iterable[ColorType]] + | NDArray[Incomplete] + | pd.Index[Incomplete] + | pd.Series[Incomplete] + | pd.DataFrame +) __all__ = ["heatmap", "clustermap"] def heatmap( - data: DataFrame | _ArrayLike[Incomplete], + data: pd.DataFrame | ArrayLike, *, vmin: float | None = None, vmax: float | None = None, @@ -38,24 +53,37 @@ def heatmap( square: bool = False, xticklabels: Literal["auto"] | bool | int | Sequence[str] = "auto", yticklabels: Literal["auto"] | bool | int | Sequence[str] = "auto", - mask: NDArray[np.bool_] | DataFrame | None = None, + mask: NDArray[np.bool_] | pd.DataFrame | None = None, ax: Axes | None = None, + # Kwargs below passed to matplotlib.axes.Axes.pcolormesh + alpha: float | None = None, + norm: str | Normalize | None = None, + shading: Literal["flat", "nearest", "gouraud", "auto"] | None = None, + antialiased: bool = False, **kwargs, ) -> Axes: ... +@type_check_only +class _Dendogram(TypedDict): + icoord: list[list[float]] + dcoord: list[list[float]] + ivl: list[str] + leaves: list[int] + color_list: list[str] + leaves_color_list: list[str] class _DendrogramPlotter: axis: int - array: NDArray[Incomplete] - data: DataFrame + array: NDArray[np.floating] + data: pd.DataFrame shape: tuple[int, int] metric: str method: str label: bool rotate: bool - linkage: NDArray[Incomplete] - dendrogram: dict[str, list[Incomplete]] - xticks: list[float] | NDArray[Incomplete] - yticks: list[float] | NDArray[Incomplete] + linkage: NDArray[np.floating] + dendrogram: _Dendogram + xticks: list[float] | NDArray[np.floating] + yticks: list[float] | NDArray[np.floating] xticklabels: list[str] yticklabels: list[str] ylabel: str @@ -63,19 +91,26 @@ class _DendrogramPlotter: dependent_coord: list[list[float]] independent_coord: list[list[float]] def __init__( - self, data: DataFrame, linkage: NDArray[Incomplete] | None, metric: str, method: str, axis: int, label: bool, rotate: bool + self, + data: pd.DataFrame, + linkage: NDArray[np.floating] | None, + metric: str, + method: str, + axis: int, + label: bool, + rotate: bool, ) -> None: ... @property - def calculated_linkage(self) -> NDArray[Incomplete]: ... - def calculate_dendrogram(self) -> dict[str, list[Incomplete]]: ... + def calculated_linkage(self) -> NDArray[np.float64]: ... + def calculate_dendrogram(self) -> _Dendogram: ... @property def reordered_ind(self) -> list[int]: ... def plot(self, ax: Axes, tree_kws: dict[str, Incomplete]) -> Self: ... def dendrogram( - data: DataFrame, + data: pd.DataFrame, *, - linkage: NDArray[Incomplete] | None = None, + linkage: NDArray[np.floating] | None = None, axis: int = 1, label: bool = True, metric: str = "euclidean", @@ -86,13 +121,13 @@ def dendrogram( ) -> _DendrogramPlotter: ... class ClusterGrid(Grid): - data: DataFrame - data2d: DataFrame - mask: DataFrame - row_colors: Incomplete - row_color_labels: Incomplete - col_colors: Incomplete - col_color_labels: Incomplete + data: pd.DataFrame + data2d: pd.DataFrame + mask: pd.DataFrame + row_colors: list[list[tuple[float, float, float]]] | None + row_color_labels: list[str] | None + col_colors: list[list[tuple[float, float, float]]] | None + col_color_labels: list[str] | None gs: GridSpec ax_row_dendrogram: Axes ax_col_dendrogram: Axes @@ -101,44 +136,38 @@ class ClusterGrid(Grid): ax_heatmap: Axes ax_cbar: Axes | None cax: Axes | None - cbar_pos: Incomplete + cbar_pos: tuple[float, float, float, float] | None dendrogram_row: _DendrogramPlotter | None dendrogram_col: _DendrogramPlotter | None def __init__( self, - data: ( - _ListLikeU - | DataFrame - | dict[Incomplete, Incomplete] - | Iterable[_ListLikeU | tuple[Hashable, _ListLikeU] | dict[Incomplete, Incomplete]] - | None - ), + data: _ConvertibleToDataFrame, pivot_kws: Mapping[str, Incomplete] | None = None, z_score: int | None = None, standard_scale: int | None = None, figsize: tuple[float, float] | None = None, - row_colors=None, - col_colors=None, - mask: NDArray[np.bool_] | DataFrame | None = None, + row_colors: _FlatOrNestedSequenceOfColors | None = None, + col_colors: _FlatOrNestedSequenceOfColors | None = None, + mask: NDArray[np.bool_] | pd.DataFrame | None = None, dendrogram_ratio: float | tuple[float, float] | None = None, colors_ratio: float | tuple[float, float] | None = None, cbar_pos: tuple[float, float, float, float] | None = None, ) -> None: ... def format_data( self, - data: DataFrame, + data: pd.DataFrame, pivot_kws: Mapping[str, Incomplete] | None, z_score: int | None = None, standard_scale: int | None = None, - ) -> DataFrame: ... + ) -> pd.DataFrame: ... @staticmethod - def z_score(data2d: DataFrame, axis: int = 1) -> DataFrame: ... + def z_score(data2d: pd.DataFrame, axis: int = 1) -> pd.DataFrame: ... @staticmethod - def standard_scale(data2d: DataFrame, axis: int = 1) -> DataFrame: ... - def dim_ratios(self, colors: Incomplete | None, dendrogram_ratio: float, colors_ratio: float) -> list[float]: ... + def standard_scale(data2d: pd.DataFrame, axis: int = 1) -> pd.DataFrame: ... + def dim_ratios(self, colors: ArrayLike | None, dendrogram_ratio: float, colors_ratio: float) -> list[float]: ... @staticmethod def color_list_to_matrix_and_cmap( - colors: Sequence[ColorType], ind: _ArrayLikeInt_co, axis: int = 0 + colors: _FlatOrNestedSequenceOfColors, ind: _ArrayLikeInt_co, axis: int = 0 ) -> tuple[NDArray[np.int_], ListedColormap]: ... def plot_dendrograms( self, @@ -146,8 +175,8 @@ class ClusterGrid(Grid): col_cluster: bool, metric: str, method: str, - row_linkage: NDArray[Incomplete] | None, - col_linkage: NDArray[Incomplete] | None, + row_linkage: NDArray[np.floating] | None, + col_linkage: NDArray[np.floating] | None, tree_kws: dict[str, Incomplete] | None, ) -> None: ... def plot_colors(self, xind: _ArrayLikeInt_co, yind: _ArrayLikeInt_co, **kws) -> None: ... @@ -159,20 +188,14 @@ class ClusterGrid(Grid): colorbar_kws: dict[str, Incomplete] | None, row_cluster: bool, col_cluster: bool, - row_linkage: NDArray[Incomplete] | None, - col_linkage: NDArray[Incomplete] | None, + row_linkage: NDArray[np.floating] | None, + col_linkage: NDArray[np.floating] | None, tree_kws: dict[str, Incomplete] | None, **kws, ) -> Self: ... def clustermap( - data: ( - _ListLikeU - | DataFrame - | dict[Incomplete, Incomplete] - | Iterable[_ListLikeU | tuple[Hashable, _ListLikeU] | dict[Incomplete, Incomplete]] - | None - ), + data: _ConvertibleToDataFrame, *, pivot_kws: dict[str, Incomplete] | None = None, method: str = "average", @@ -183,11 +206,11 @@ def clustermap( cbar_kws: dict[str, Incomplete] | None = None, row_cluster: bool = True, col_cluster: bool = True, - row_linkage: NDArray[Incomplete] | None = None, - col_linkage: NDArray[Incomplete] | None = None, - row_colors=None, - col_colors=None, - mask: NDArray[np.bool_] | DataFrame | None = None, + row_linkage: NDArray[np.floating] | None = None, + col_linkage: NDArray[np.floating] | None = None, + row_colors: _FlatOrNestedSequenceOfColors | None = None, + col_colors: _FlatOrNestedSequenceOfColors | None = None, + mask: NDArray[np.bool_] | pd.DataFrame | None = None, dendrogram_ratio: float | tuple[float, float] = 0.2, colors_ratio: float | tuple[float, float] = 0.03, cbar_pos: tuple[float, float, float, float] | None = (0.02, 0.8, 0.05, 0.18),