From 3fe9674712dbc8adb2489fe685f3f1d584c2b42b Mon Sep 17 00:00:00 2001 From: Jan-Eric Nitschke <47750513+JanEricNitschke@users.noreply.github.com> Date: Sat, 13 Jul 2024 11:23:50 +0200 Subject: [PATCH 1/3] Fix series.map overloads --- pandas-stubs/core/series.pyi | 2 +- tests/test_series.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 209237746..2b163e138 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -917,7 +917,7 @@ class Series(IndexOpsMixin[S1], NDFrame): def map( self, arg: Callable[[S1], S2 | NAType] | Mapping[S1, S2] | Series[S2], - na_action: Literal["ignore"] = ..., + na_action: Literal["ignore"], ) -> Series[S2]: ... @overload def map( diff --git a/tests/test_series.py b/tests/test_series.py index 9c99dba41..503bb685c 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -3276,18 +3276,34 @@ def test_map_na() -> None: mapping = {1: "a", 2: "b", 3: "c"} check(assert_type(s.map(mapping, na_action=None), "pd.Series[str]"), pd.Series, str) + check(assert_type(s.map(mapping), "pd.Series[str]"), pd.Series, str) def callable(x: int | NAType) -> str | NAType: if isinstance(x, int): return str(x) return x + def bad_callable(x: int) -> int: + return x + 1 + + s.map( + bad_callable, na_action=None # type: ignore[arg-type] # pyright: ignore[reportCallIssue, reportArgumentType] + ) + 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 ) + 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) + check(assert_type(s.map(series), "pd.Series[str]"), pd.Series, str) def test_case_when() -> None: From 7adb42d7334d010f3bbfd99416c06e02dc84e1ee Mon Sep 17 00:00:00 2001 From: Jan-Eric Nitschke <47750513+JanEricNitschke@users.noreply.github.com> Date: Mon, 15 Jul 2024 16:34:02 +0200 Subject: [PATCH 2/3] Fix tests to properly exercise lack of 'pd.NA' support. --- tests/test_series.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/test_series.py b/tests/test_series.py index 503bb685c..e723c84ed 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -3284,12 +3284,14 @@ def callable(x: int | NAType) -> str | NAType: return x def bad_callable(x: int) -> int: - return x + 1 + return x << 1 - s.map( - bad_callable, na_action=None # type: ignore[arg-type] # pyright: ignore[reportCallIssue, reportArgumentType] - ) - s.map(bad_callable) # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + 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, From 8810cc7ee06f9ad39977da0b9faf0f809f447f44 Mon Sep 17 00:00:00 2001 From: Jan-Eric Nitschke <47750513+JanEricNitschke@users.noreply.github.com> Date: Tue, 16 Jul 2024 18:20:05 +0200 Subject: [PATCH 3/3] Further splitting of map overloads --- pandas-stubs/_typing.pyi | 2 ++ pandas-stubs/core/series.pyi | 7 ++++++ tests/test_series.py | 49 +++++++++++++++++++++++++----------- 3 files changed, 43 insertions(+), 15 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index baa358a8a..7ec2fe200 100644 --- a/pandas-stubs/_typing.pyi +++ b/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 ) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 2b163e138..13a657d8d 100644 --- a/pandas-stubs/core/series.pyi +++ b/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], diff --git a/tests/test_series.py b/tests/test_series.py index e723c84ed..306910323 100644 --- a/tests/test_series.py +++ b/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")