Skip to content

Commit

Permalink
Add stub for Styler.map_index (#905)
Browse files Browse the repository at this point in the history
* fix: add stub for Styler.map_index

* add test for Styler.map_index

* Fix map_index type definition

* Update test_map_index to show example from docs

* Run pre-commit
  • Loading branch information
cquick01 committed Apr 10, 2024
1 parent 9dd9eee commit b30d777
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 0 deletions.
7 changes: 7 additions & 0 deletions pandas-stubs/io/formats/style.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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: ...
Expand Down
9 changes: 9 additions & 0 deletions tests/test_styler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down

0 comments on commit b30d777

Please sign in to comment.