diff --git a/pandas-stubs/io/formats/style.pyi b/pandas-stubs/io/formats/style.pyi index 240a161ff..dc2b53829 100644 --- a/pandas-stubs/io/formats/style.pyi +++ b/pandas-stubs/io/formats/style.pyi @@ -218,6 +218,13 @@ class Styler(StylerRenderer): level: Level | list[Level] | None = ..., **kwargs: Any, ) -> Styler: ... + def map_index( + self, + func: Callable[[Scalar], str | None], + axis: Axis = ..., + level: Level | list[Level] | None = ..., + **kwargs, + ) -> Styler: ... def set_table_attributes(self, attributes: str) -> Styler: ... def export(self) -> StyleExportDict: ... def use(self, styles: StyleExportDict) -> Styler: ... diff --git a/tests/test_styler.py b/tests/test_styler.py index 188710787..7dd4003ec 100644 --- a/tests/test_styler.py +++ b/tests/test_styler.py @@ -19,6 +19,8 @@ import pytest from typing_extensions import assert_type +from pandas._typing import Scalar + from tests import check from pandas.io.formats.style import Styler @@ -63,6 +65,13 @@ def f1(s: Series) -> Series[str]: check(assert_type(DF.style.apply_index(f1), 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 test_background_gradient() -> None: check(assert_type(DF.style.background_gradient(), Styler), Styler)