Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 6 additions & 15 deletions pandas-stubs/core/series.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -1103,25 +1103,16 @@ 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,
arg: (
Callable[[S1], S2 | None | NAType]
| Mapping[S1, S2 | None | NAType]
| Series[S2]
),
na_action: Literal["ignore"] | None = None,
) -> Series[S2]: ...
@overload
def map(
self,
arg: Callable[[Any], Any] | Mapping[Any, Any] | Series,
na_action: Literal["ignore"] | None = ...,
) -> Series: ...
@overload
def aggregate(
self: Series[int],
func: Literal["mean"],
Expand Down
28 changes: 27 additions & 1 deletion tests/series/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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,
)


Expand All @@ -3620,10 +3634,22 @@ def callable(x: int | NAType) -> str | NAType:
check(
assert_type(s.map(callable, na_action=None), "pd.Series[str]"), pd.Series, str
)
# na_action defaults to None
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=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")
Expand Down