Fix signature of Series.map() (#942)

* Fix signature of Series.map()

* Add tests for Series.map() hints.

Author: JanEricNitschke

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
 )

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],

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)