styler

Author: cmp0xff

pandas-stubs/io/formats/style.pyi

@@ -8,7 +8,6 @@ from typing import (
     Any,
     Concatenate,
     Literal,
-    Protocol,
     Self,
     overload,
 )
@@ -52,26 +51,6 @@ from pandas.io.formats.style_render import (
     Subset,
 )
 
-class _SeriesFunc(Protocol):
-    def __call__(
-        self, series: Series, /, *args: Any, **kwargs: Any
-    ) -> list[Any] | Series: ...
-
-class _SeriesStrFunc(Protocol):
-    def __call__(
-        self, series: Series[str], /, *args: Any, **kwargs: Any
-    ) -> list[str] | Series[str]: ...
-
-class _DataFrameFunc(Protocol):
-    def __call__(
-        self, series: DataFrame, /, *args: Any, **kwargs: Any
-    ) -> np_ndarray | DataFrame: ...
-
-class _MapCallable(Protocol):
-    def __call__(
-        self, first_arg: Scalar, /, *args: Any, **kwargs: Any
-    ) -> str | None: ...
-
 class Styler(StylerRenderer):
     def __init__(
         self,
@@ -90,16 +69,12 @@ class Styler(StylerRenderer):
         formatter: ExtFormatter | None = ...,
     ) -> None: ...
     def concat(self, other: Styler) -> Styler: ...
-    @overload
-    def map(
-        self,
-        func: Callable[[Scalar], str | None],
-        subset: Subset[Hashable] | None = ...,
-    ) -> Styler: ...
-    @overload
     def map(
         self,
-        func: _MapCallable,
+        func: (
+            Callable[[Scalar], str | None]
+            | Callable[Concatenate[Scalar, ...], str | None]
+        ),
         subset: Subset[Hashable] | None = ...,
         **kwargs: Any,
     ) -> Styler: ...
@@ -267,32 +242,43 @@ class Styler(StylerRenderer):
     @overload
     def apply(
         self,
-        func: _SeriesFunc | Callable[[Series], list[Any] | Series],
+        func: (
+            Callable[[Series], list[Any] | Series]
+            | Callable[Concatenate[Series, ...], list[Any] | Series]
+        ),
         axis: Axis = ...,
         subset: Subset[Hashable] | None = ...,
         **kwargs: Any,
     ) -> Styler: ...
     @overload
     def apply(
         self,
-        func: _DataFrameFunc | Callable[[DataFrame], np_ndarray | DataFrame],
+        func: (
+            Callable[[DataFrame], np_ndarray | DataFrame]
+            | Callable[Concatenate[DataFrame, ...], np_ndarray | DataFrame]
+        ),
         axis: None,
         subset: Subset[Hashable] | None = ...,
         **kwargs: Any,
     ) -> Styler: ...
     def apply_index(
         self,
         func: (
-            _SeriesStrFunc
-            | Callable[[Series], list[str] | np_ndarray_str | Series[str]]
+            Callable[[Series], list[str] | np_ndarray_str | Series[str]]
+            | Callable[
+                Concatenate[Series, ...], list[str] | np_ndarray_str | Series[str]
+            ]
         ),
         axis: Axis = ...,
         level: Level | list[Level] | None = ...,
         **kwargs: Any,
     ) -> Styler: ...
     def map_index(
         self,
-        func: _MapCallable | Callable[[Scalar], str | None],
+        func: (
+            Callable[[Scalar], str | None]
+            | Callable[Concatenate[Scalar, ...], str | None]
+        ),
         axis: Axis = ...,
         level: Level | list[Level] | None = ...,
         **kwargs: Any,

tests/test_styler.py

@@ -60,10 +60,7 @@ def highlight_max(x: Series[int], /, color: str) -> list[str]:
         return [f"color: {color}" if val == x.max() else "" for val in x]
 
     check(
-        # TODO: facebook/pyrefly#0
-        # pyrefly: ignore[assert-type,no-matching-overload]
-        assert_type(DF.style.apply(highlight_max, color="red", axis=1), Styler),
-        Styler,
+        assert_type(DF.style.apply(highlight_max, color="red", axis=1), Styler), Styler
     )
 
 
@@ -84,10 +81,7 @@ def highlight_odd(index: pd.Series, color: str) -> list[str]:
 
     check(
         assert_type(
-            # TODO: facebook/pyrefly#0
-            # pyrefly: ignore[bad-argument-type]
-            DF.style.apply_index(highlight_odd, axis=0, color="purple"),
-            Styler,
+            DF.style.apply_index(highlight_odd, axis=0, color="purple"), Styler
         ),
         Styler,
     )
@@ -102,8 +96,6 @@ def f(s: Scalar) -> str | None:
     def f1(s: Scalar, color: str) -> str | None:
         return f"background-color: {color};" if s == "b" else None
 
-    # TODO: facebook/pyrefly#0
-    # pyrefly: ignore[bad-argument-type]
     check(assert_type(DF.style.map_index(f1, color="pink", axis=0), Styler), Styler)
 
 
@@ -286,8 +278,6 @@ def color_negative(v: Scalar, /, color: str) -> str | None:
 
     df = DataFrame(np.random.randn(5, 2), columns=["A", "B"])
 
-    # TODO: facebook/pyrefly#0
-    # pyrefly: ignore[assert-type,no-matching-overload]
     check(assert_type(df.style.map(color_negative, color="red"), Styler), Styler)