Further splitting of map overloads

Author: JanEricNitschke

pandas-stubs/_typing.pyi

@@ -568,6 +568,8 @@ S2 = TypeVar(
     | BaseOffset,
 )
 
+SN = TypeVar("SN", bound=bool | int | float | complex)
+
 IndexingInt: TypeAlias = (
     int | np.int_ | np.integer | np.unsignedinteger | np.signedinteger | np.int8
 )

pandas-stubs/core/series.pyi

@@ -94,6 +94,7 @@ from pandas._libs.tslibs.nattype import NaTType
 from pandas._typing import (
     S1,
     S2,
+    SN,
     AggFuncTypeBase,
     AggFuncTypeDictFrame,
     AggFuncTypeSeriesToFrame,
@@ -914,6 +915,12 @@ class Series(IndexOpsMixin[S1], NDFrame):
         fill_value: int | _str | dict | None = ...,
     ) -> DataFrame: ...
     @overload
+    def map(
+        self,
+        arg: Callable[[SN], S2 | NAType] | Mapping[SN, S2] | Series[S2],
+        na_action: Literal["ignore"] = ...,
+    ) -> Series[S2]: ...
+    @overload
     def map(
         self,
         arg: Callable[[S1], S2 | NAType] | Mapping[S1, S2] | Series[S2],

tests/test_series.py

@@ -3283,21 +3283,6 @@ def callable(x: int | NAType) -> str | NAType:
             return str(x)
         return x
 
-    def bad_callable(x: int) -> int:
-        return x << 1
-
-    with pytest.raises(TypeError):
-        s.map(
-            bad_callable, na_action=None  # type: ignore[arg-type] # pyright: ignore[reportCallIssue, reportArgumentType]
-        )
-    with pytest.raises(TypeError):
-        s.map(bad_callable)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
-    check(
-        assert_type(s.map(bad_callable, na_action="ignore"), "pd.Series[int]"),
-        pd.Series,
-        int,
-    )
-
     check(
         assert_type(s.map(callable, na_action=None), "pd.Series[str]"), pd.Series, str
     )
@@ -3307,6 +3292,40 @@ def bad_callable(x: int) -> int:
     check(assert_type(s.map(series, na_action=None), "pd.Series[str]"), pd.Series, str)
     check(assert_type(s.map(series), "pd.Series[str]"), pd.Series, str)
 
+    s2: pd.Series[float] = pd.Series([1.0, pd.NA, 3.0])
+
+    def callable2(x: float) -> float:
+        return x + 1
+
+    check(
+        assert_type(s2.map(callable2, na_action="ignore"), "pd.Series[float]"),
+        pd.Series,
+        float,
+    )
+    check(
+        assert_type(s2.map(callable2), "pd.Series[float]"),
+        pd.Series,
+        float,
+    )
+    if TYPE_CHECKING_INVALID_USAGE:
+        s2.map(callable2, na_action=None)  # type: ignore[arg-type] # pyright: ignore[reportCallIssue, reportArgumentType]
+
+    s3: pd.Series[str] = pd.Series(["A", pd.NA, "C"])
+
+    def callable3(x: str) -> str:
+        return x.lower()
+
+    check(
+        assert_type(s3.map(callable3, na_action="ignore"), "pd.Series[str]"),
+        pd.Series,
+        str,
+    )
+    if TYPE_CHECKING_INVALID_USAGE:
+        s3.map(
+            callable3, na_action=None  # type: ignore[arg-type] # pyright: ignore[reportCallIssue, reportArgumentType]
+        )
+        s3.map(callable3)  # type: ignore[type-var] # pyright: ignore[reportCallIssue, reportArgumentType]
+
 
 def test_case_when() -> None:
     c = pd.Series([6, 7, 8, 9], name="c")