Simplify and add more tests

Author: hamdanal

pandas-stubs/core/series.pyi

@@ -1103,24 +1103,15 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
         fill_value: int | _str | dict | None = None,
         sort: _bool = True,
     ) -> DataFrame: ...
-    @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 = None,
-    ) -> Series[S2]: ...
-    @overload
     def map(
         self,
-        arg: Callable[[Any], Any] | Mapping[Any, Any] | Series,
+        arg: (
+            Callable[[S1], S2 | None | NAType]
+            | Mapping[S1, S2 | None | NAType]
+            | Series[S2]
+        ),
         na_action: Literal["ignore"] | None = None,
-    ) -> Series: ...
+    ) -> Series[S2]: ...
     @overload
     def aggregate(
         self: Series[int],

tests/series/test_series.py

@@ -3584,6 +3584,7 @@ def test_map() -> None:
         pd.Series,
         str,
     )
+    check(assert_type(s.map(mapping), "pd.Series[str]"), pd.Series, str)
 
     def callable(x: int) -> str:
         return str(x)
@@ -3593,16 +3594,29 @@ def callable(x: int) -> str:
         pd.Series,
         str,
     )
+    check(assert_type(s.map(callable), "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
     )
+    check(assert_type(s.map(series), "pd.Series[str]"), pd.Series, str)
 
     unknown_series = pd.Series([1, 0, None])
     check(
-        assert_type(unknown_series.map({1: True, 0: False, None: None}), pd.Series),
+        assert_type(
+            unknown_series.map({1: True, 0: False, None: None}), "pd.Series[bool]"
+        ),
+        pd.Series,
+        bool,
+    )
+    check(
+        assert_type(
+            unknown_series.map({1: True, 0: False, None: None}, na_action="ignore"),
+            "pd.Series[bool]",
+        ),
         pd.Series,
+        bool,
     )
 
 
@@ -3626,6 +3640,16 @@ def callable(x: int | NAType) -> str | NAType:
     series = pd.Series(["a", "b", "c"])
     check(assert_type(s.map(series, na_action=None), "pd.Series[str]"), pd.Series, str)
 
+    def callable2(x: int | NAType | None) -> str | None:
+        if isinstance(x, int):
+            return str(x)
+        return None
+
+    check(
+        assert_type(s.map(callable2, na_action=None), "pd.Series[str]"), pd.Series, str
+    )
+    check(assert_type(s.map(callable2), "pd.Series[str]"), pd.Series, str)
+
 
 def test_case_when() -> None:
     c = pd.Series([6, 7, 8, 9], name="c")