Skip to content

GH1226 Add Styler.map to the stubs #1228

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
May 23, 2025
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions pandas-stubs/io/formats/style.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,19 @@ 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 = ...,
) -> Styler: ...
@overload
def map(
self,
func: Callable[..., str | None],
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Dr-Irv I could not find how to type a function that takes Scalar and some other stuff, I tried [Scalar, ...] but it was not happy, feel free to recommend something better.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Callable cannot express complex signatures

I think this is usually done using Protocol, I tried writing this:

    class _MapCallable(Protocol):
        def __call__(self, first_arg: Scalar, /, *args: Any, **kwargs: Any) -> str | None: ...

    def map(
        self,
        func: _MapCallable,
        subset: Subset | None = ...,
        **kwargs: Any,
    ) -> Styler: ...

but it fails the test

tests/test_styler.py:256: error: Expression is of type "Any", not "Styler"  [assert-type]
tests/test_styler.py:257: error: No overload variant of "map" of "Styler" matches argument types "Callable[[str | bytes | date | datetime | timedelta | <7 more items> | complex | integer[Any] | floating[Any] | complexfloating[Any, Any], str], str | None]", "str"  [call-overload]
tests/test_styler.py:257: note: Possible overload variants:
tests/test_styler.py:257: note:     def map(self, func: Callable[[str | bytes | date | datetime | timedelta | <7 more items> | complex | integer[Any] | floating[Any] | complexfloating[Any, Any]], str | None], subset: _IndexSlice | slice[Any, Any, Any] | tuple[slice[Any, Any, Any], ...] | list[Any] | Index[Any] | None = ...) -> Styler
tests/test_styler.py:257: note:     def map(self, func: _MapCallable, subset: _IndexSlice | slice[Any, Any, Any] | tuple[slice[Any, Any, Any], ...] | list[Any] | Index[Any] | None = ..., **kwargs: Any) -> Styler

Maybe this will be helpful to figure something out.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The idea of using the Protocol works if we change the test. Instead of

    def color_negative(v: Scalar,  color: str) -> str | None:

if we use

    def color_negative(v: Scalar,  /, color: str) -> str | None:

then it is fine. Note that pyright accepts the func without the / defined, but mypy doesn't, so this is a workaround to make mypy happy.

@loicdiridollou change the test that way and put the _MapCallable near the other Protocol definition in that file and we should be good to go.

@Molkree Thanks for the suggestion!

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that pyright accepts the func without the / defined, but mypy doesn't, so this is a workaround to make mypy happy.

I had to double check that and I agree that this is mypy bug. According to these typing rules callable with standard argument of matching type should be assignable to the positional-only argument of matching type, same for kwargs but because we typed them as Any everything should be assignable to everything.

I think it's this bug specifically: python/mypy#12525 (or a variation of it, error messages are not the same)

You learn something new everyday!

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for all the feedback here, indeed I learn about those intricacies and mypy issues!
My only fear is that it would force people to have the argument-based function with the /, but I think it is worth trying out and seeing if we get complaints from the user base.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for all the feedback here, indeed I learn about those intricacies and mypy issues! My only fear is that it would force people to have the argument-based function with the /, but I think it is worth trying out and seeing if we get complaints from the user base.

It only forces people who use mypy to include /. If you use pyright, it isn't required.

subset: Subset | None = ...,
**kwargs: Any,
) -> Styler: ...
def set_tooltips(
self,
ttips: DataFrame,
Expand Down
28 changes: 27 additions & 1 deletion tests/test_styler.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@

DF = DataFrame({"a": [1, 2, 3], "b": [3.14, 2.72, 1.61]})


PWD = pathlib.Path(os.path.split(os.path.abspath(__file__))[0])

if TYPE_CHECKING:
Expand Down Expand Up @@ -233,3 +232,30 @@ def test_styler_columns_and_index() -> None:
styler = DF.style
check(assert_type(styler.columns, Index), Index)
check(assert_type(styler.index, Index), Index)


def test_styler_map() -> None:
"""Test type returned with Styler.map GH1226."""
df = DataFrame(data={"col1": [1, -2], "col2": [-3, 4]})
check(
assert_type(
df.style.map(
lambda v: "color: red;" if isinstance(v, float) and v < 0 else None
),
Styler,
),
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,
)