From b30d777fb219ce5da0ffb374a835a536d832cfb8 Mon Sep 17 00:00:00 2001
From: cquick01 <72400253+cquick01@users.noreply.github.com>
Date: Wed, 10 Apr 2024 13:19:55 -0400
Subject: [PATCH] Add stub for Styler.map_index (#905)

* 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
---
 pandas-stubs/io/formats/style.pyi | 7 +++++++
 tests/test_styler.py              | 9 +++++++++
 2 files changed, 16 insertions(+)

diff --git a/pandas-stubs/io/formats/style.pyi b/pandas-stubs/io/formats/style.pyi
index 240a161f..dc2b5382 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 18871078..7dd4003e 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)