GH1226 PR Feedback

Author: loicdiridollou

pandas-stubs/io/formats/style.pyi

@@ -71,11 +71,18 @@ class Styler(StylerRenderer):
         formatter: ExtFormatter | None = ...,
     ) -> None: ...
     def concat(self, other: Styler) -> Styler: ...
+    @overload
     def map(
         self,
         func: Callable[[Scalar], str | None],
         subset: Subset | None = ...,
-        **kwargs: dict[str, Any],
+    ) -> Styler: ...
+    @overload
+    def map(
+        self,
+        func: Callable[..., str | None],
+        subset: Subset | None = ...,
+        **kwargs: Any,
     ) -> Styler: ...
     def set_tooltips(
         self,

tests/test_styler.py

@@ -2,7 +2,10 @@
 
 import os
 import pathlib
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
 
 from jinja2.environment import (
     Environment,
@@ -246,3 +249,16 @@ def test_styler_map() -> None:
         ),
         Styler,
     )
+
+    def color_negative(v: Scalar, color: str) -> str | None:
+        return f"color: {color};" if isinstance(v, float) and v < 0 else None
+
+    df = DataFrame(np.random.randn(5, 2), columns=["A", "B"])
+
+    check(
+        assert_type(
+            df.style.map(color_negative, color="red"),
+            Styler,
+        ),
+        Styler,
+    )