seaborn: improve the matrix module (#14464)

Author: hamdanal

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,44 +53,64 @@ 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
     xlabel: str
     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,53 +136,47 @@ 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,
         row_cluster: bool,
         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),