use protocol for styler.apply

Author: Dr-Irv

pandas-stubs/io/formats/style.pyi

@@ -42,7 +42,14 @@ from pandas.io.formats.style_render import (
 )
 
 class SeriesFunc(Protocol):
-    def __call__(self, series: Series, *args: Any, **kwargs: Any) -> list | Series: ...
+    def __call__(
+        self, series: Series, /, *args: Any, **kwargs: Any
+    ) -> list | Series: ...
+
+class DataFrameFunc(Protocol):
+    def __call__(
+        self, series: DataFrame, /, *args: Any, **kwargs: Any
+    ) -> npt.NDArray | DataFrame: ...
 
 class Styler(StylerRenderer):
     def __init__(
@@ -202,15 +209,15 @@ class Styler(StylerRenderer):
     @overload
     def apply(
         self,
-        func: SeriesFunc,
+        func: SeriesFunc | Callable[[Series], list | Series],
         axis: Axis = ...,
         subset: Subset | None = ...,
         **kwargs: Any,
     ) -> Styler: ...
     @overload
     def apply(
         self,
-        func: Callable[[DataFrame], npt.NDArray | DataFrame],
+        func: DataFrameFunc | Callable[[DataFrame], npt.NDArray | DataFrame],
         axis: None,
         subset: Subset | None = ...,
         **kwargs: Any,

tests/test_styler.py

@@ -52,6 +52,14 @@ def h(df: DataFrame) -> DataFrame:
 
     check(assert_type(DF.style.apply(h, axis=None), Styler), Styler)
 
+    # GH 919
+    def highlight_max(x: Series[int], /, color: str) -> list[str]:
+        return [f"color: {color}" if val == x.max() else "" for val in x]
+
+    check(
+        assert_type(DF.style.apply(highlight_max, color="red", axis=1), Styler), Styler
+    )
+
 
 def test_apply_index() -> None:
     def f(s: Series) -> npt.NDArray[np.str_]: