Skip to content

Commit d263eea

Browse files
Add tests for Series.map() hints.
1 parent bfce13e commit d263eea

File tree

1 file changed

+44
-0
lines changed

1 file changed

+44
-0
lines changed

tests/test_series.py

+44
Original file line numberDiff line numberDiff line change
@@ -3239,3 +3239,47 @@ def test_operator_constistency() -> None:
32393239
pd.Series,
32403240
pd.Timedelta,
32413241
)
3242+
3243+
3244+
def test_map() -> None:
3245+
s = pd.Series([1, 2, 3])
3246+
3247+
mapping = {1: "a", 2: "b", 3: "c"}
3248+
check(
3249+
assert_type(s.map(mapping, na_action="ignore"), "pd.Series[str]"),
3250+
pd.Series,
3251+
str,
3252+
)
3253+
3254+
def callable(x: int) -> str:
3255+
return str(x)
3256+
3257+
check(
3258+
assert_type(s.map(callable, na_action="ignore"), "pd.Series[str]"),
3259+
pd.Series,
3260+
str,
3261+
)
3262+
3263+
series = pd.Series(["a", "b", "c"])
3264+
check(
3265+
assert_type(s.map(series, na_action="ignore"), "pd.Series[str]"), pd.Series, str
3266+
)
3267+
3268+
3269+
def test_map_na() -> None:
3270+
s: pd.Series[int] = pd.Series([1, pd.NA, 3])
3271+
3272+
mapping = {1: "a", 2: "b", 3: "c"}
3273+
check(assert_type(s.map(mapping, na_action=None), "pd.Series[str]"), pd.Series, str)
3274+
3275+
def callable(x: int | NAType) -> str | NAType:
3276+
if isinstance(x, int):
3277+
return str(x)
3278+
return x
3279+
3280+
check(
3281+
assert_type(s.map(callable, na_action=None), "pd.Series[str]"), pd.Series, str
3282+
)
3283+
3284+
series = pd.Series(["a", "b", "c"])
3285+
check(assert_type(s.map(series, na_action=None), "pd.Series[str]"), pd.Series, str)

0 commit comments

Comments
 (0)