From 411f2dea275d99dadecc05ac2255eae741a3d444 Mon Sep 17 00:00:00 2001
From: Jan-Eric Nitschke <47750513+JanEricNitschke@users.noreply.github.com>
Date: Mon, 24 Jun 2024 22:10:39 +0200
Subject: [PATCH 1/2] Fix signature of Series.map()

---
 pandas-stubs/_typing.pyi     | 19 +++++++++++++++++++
 pandas-stubs/core/series.pyi | 14 +++++++++++++-
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi
index 28dc09512..f4cc75370 100644
--- a/pandas-stubs/_typing.pyi
+++ b/pandas-stubs/_typing.pyi
@@ -549,6 +549,25 @@ S1 = TypeVar(
     | BaseOffset,
 )
 
+S2 = TypeVar(
+    "S2",
+    bound=str
+    | bytes
+    | datetime.date
+    | datetime.time
+    | bool
+    | int
+    | float
+    | complex
+    | Dtype
+    | datetime.datetime  # includes pd.Timestamp
+    | datetime.timedelta  # includes pd.Timedelta
+    | Period
+    | Interval
+    | CategoricalDtype
+    | BaseOffset,
+)
+
 IndexingInt: TypeAlias = (
     int | np.int_ | np.integer | np.unsignedinteger | np.signedinteger | np.int8
 )
diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
index 27fd6c6bc..d6ccca521 100644
--- a/pandas-stubs/core/series.pyi
+++ b/pandas-stubs/core/series.pyi
@@ -93,6 +93,7 @@ from pandas._libs.tslibs import BaseOffset
 from pandas._libs.tslibs.nattype import NaTType
 from pandas._typing import (
     S1,
+    S2,
     AggFuncTypeBase,
     AggFuncTypeDictFrame,
     AggFuncTypeSeriesToFrame,
@@ -913,7 +914,18 @@ class Series(IndexOpsMixin[S1], NDFrame):
         level: Level = ...,
         fill_value: int | _str | dict | None = ...,
     ) -> DataFrame: ...
-    def map(self, arg, na_action: Literal["ignore"] | None = ...) -> Series[S1]: ...
+    @overload
+    def map(
+        self,
+        arg: Callable[[S1], S2 | NAType] | Mapping[S1, S2] | Series[S2],
+        na_action: Literal["ignore"] = ...,
+    ) -> Series[S2]: ...
+    @overload
+    def map(
+        self,
+        arg: Callable[[S1 | NAType], S2 | NAType] | Mapping[S1, S2] | Series[S2],
+        na_action: None = ...,
+    ) -> Series[S2]: ...
     @overload
     def aggregate(  # type: ignore[overload-overlap]
         self: Series[int],

From 2c9464d911ecfa5847057582c20d8572a8ba76b1 Mon Sep 17 00:00:00 2001
From: Jan-Eric Nitschke <47750513+JanEricNitschke@users.noreply.github.com>
Date: Tue, 25 Jun 2024 18:00:19 +0200
Subject: [PATCH 2/2] Add tests for Series.map() hints.

---
 tests/test_series.py | 44 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/tests/test_series.py b/tests/test_series.py
index 61e2b8bbb..35e9ee063 100644
--- a/tests/test_series.py
+++ b/tests/test_series.py
@@ -3230,3 +3230,47 @@ def test_operator_constistency() -> None:
         pd.Series,
         pd.Timedelta,
     )
+
+
+def test_map() -> None:
+    s = pd.Series([1, 2, 3])
+
+    mapping = {1: "a", 2: "b", 3: "c"}
+    check(
+        assert_type(s.map(mapping, na_action="ignore"), "pd.Series[str]"),
+        pd.Series,
+        str,
+    )
+
+    def callable(x: int) -> str:
+        return str(x)
+
+    check(
+        assert_type(s.map(callable, na_action="ignore"), "pd.Series[str]"),
+        pd.Series,
+        str,
+    )
+
+    series = pd.Series(["a", "b", "c"])
+    check(
+        assert_type(s.map(series, na_action="ignore"), "pd.Series[str]"), pd.Series, str
+    )
+
+
+def test_map_na() -> None:
+    s: pd.Series[int] = pd.Series([1, pd.NA, 3])
+
+    mapping = {1: "a", 2: "b", 3: "c"}
+    check(assert_type(s.map(mapping, na_action=None), "pd.Series[str]"), pd.Series, str)
+
+    def callable(x: int | NAType) -> str | NAType:
+        if isinstance(x, int):
+            return str(x)
+        return x
+
+    check(
+        assert_type(s.map(callable, na_action=None), "pd.Series[str]"), pd.Series, str
+    )
+
+    series = pd.Series(["a", "b", "c"])
+    check(assert_type(s.map(series, na_action=None), "pd.Series[str]"), pd.Series, str)