BUG: allow kwargs in Styler apply_index and map_index (#1725)

* BUG: allow kwargs in Styler apply_index and map_index

* fixup! BUG: allow kwargs in Styler apply_index and map_index

Author: eosti

pandas-stubs/io/formats/style.pyi

@@ -57,6 +57,11 @@ class _SeriesFunc(Protocol):
         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
@@ -277,14 +282,17 @@ class Styler(StylerRenderer):
     ) -> Styler: ...
     def apply_index(
         self,
-        func: Callable[[Series], list[str] | np_ndarray_str | Series[str]],
+        func: (
+            _SeriesStrFunc
+            | Callable[[Series], list[str] | np_ndarray_str | Series[str]]
+        ),
         axis: Axis = ...,
         level: Level | list[Level] | None = ...,
         **kwargs: Any,
     ) -> Styler: ...
     def map_index(
         self,
-        func: Callable[[Scalar], str | None],
+        func: _MapCallable | Callable[[Scalar], str | None],
         axis: Axis = ...,
         level: Level | list[Level] | None = ...,
         **kwargs: Any,

tests/test_styler.py

@@ -76,13 +76,29 @@ def f1(s: Series) -> Series[str]:
 
     check(assert_type(DF.style.apply_index(f1), Styler), Styler)
 
+    # GH 1723
+    def highlight_odd(index: pd.Series, color: str) -> list[str]:
+        return [f"color: {color}" if x % 2 else "" for x in index]
+
+    check(
+        assert_type(
+            DF.style.apply_index(highlight_odd, axis=0, color="purple"), Styler
+        ),
+        Styler,
+    )
+
 
 def test_map_index() -> None:
     def f(s: Scalar) -> str | None:
         return "background-color: yellow;" if s == "B" else None
 
     check(assert_type(DF.style.map_index(f), Styler), Styler)
 
+    def f1(s: Scalar, color: str) -> str | None:
+        return f"background-color: {color};" if s == "b" else None
+
+    check(assert_type(DF.style.map_index(f1, color="pink", axis=0), Styler), Styler)
+
 
 def test_background_gradient() -> None:
     check(assert_type(DF.style.background_gradient(), Styler), Styler)