From e85659a6a16efacc26b458ad5c0d636bdd3efd09 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Fri, 7 Feb 2025 14:46:05 +0000 Subject: [PATCH 1/6] type freq in shift --- pandas-stubs/core/frame.pyi | 3 ++- pandas-stubs/core/series.pyi | 3 ++- tests/test_frame.py | 5 ++++- tests/test_series.py | 3 ++- 4 files changed, 10 insertions(+), 4 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 9528755a2..1089e360c 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -84,6 +84,7 @@ from pandas._typing import ( FilePath, FillnaOptions, FormattersType, + Frequency, GroupByObjectNonScalar, HashableT, HashableT1, @@ -855,7 +856,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): def shift( self, periods: int = ..., - freq=..., + freq: Frequency | dt.timedelta | None = ..., axis: Axis = ..., fill_value: Hashable | None = ..., ) -> Self: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 8b78c6d16..d7f68c820 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -119,6 +119,7 @@ from pandas._typing import ( FilePath, FillnaOptions, FloatDtypeArg, + Frequency, GroupByObjectNonScalar, HashableT1, IgnoreRaise, @@ -1219,7 +1220,7 @@ class Series(IndexOpsMixin[S1], NDFrame): def shift( self, periods: int = ..., - freq=..., + freq: Frequency | timedelta | None = ..., axis: AxisIndex = ..., fill_value: object | None = ..., ) -> Series[S1]: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index fe2b6be1b..f93c59a20 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -516,10 +516,13 @@ def test_types_sort_values_with_key() -> None: def test_types_shift() -> None: - df = pd.DataFrame(data={"col1": [1, 1], "col2": [3, 4]}) + df = pd.DataFrame( + data={"col1": [1, 1], "col2": [3, 4]}, index=pd.date_range("2020", periods=2) + ) df.shift() df.shift(1) df.shift(-1) + df.shift(freq="1D") def test_types_rank() -> None: diff --git a/tests/test_series.py b/tests/test_series.py index f0dbe8019..91cf757e5 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -416,10 +416,11 @@ def test_types_sort_values_with_key() -> None: def test_types_shift() -> None: - s = pd.Series([1, 2, 3]) + s = pd.Series([1, 2, 3], index=pd.date_range("2020", periods=3)) s.shift() s.shift(axis=0, periods=1) s.shift(-1, fill_value=0) + s.shift(freq="1D") def test_types_rank() -> None: From e4e21b5f68a89aa77fe62cd8ff6dabe021a1427f Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Sat, 8 Feb 2025 12:14:23 +0000 Subject: [PATCH 2/6] assert return type --- tests/test_series.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/test_series.py b/tests/test_series.py index 91cf757e5..90f40537b 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -417,10 +417,19 @@ def test_types_sort_values_with_key() -> None: def test_types_shift() -> None: s = pd.Series([1, 2, 3], index=pd.date_range("2020", periods=3)) - s.shift() - s.shift(axis=0, periods=1) - s.shift(-1, fill_value=0) - s.shift(freq="1D") + # Return type "Series[int]"" not quite correct + # - https://github.com/pandas-dev/pandas-stubs/issues/1111 + # - https://github.com/pandas-dev/pandas-stubs/issues/1110 + check(assert_type(s.shift(), "pd.Series[int]"), pd.Series, np.floating) + check( + assert_type(s.shift(axis=0, periods=1), "pd.Series[int]"), + pd.Series, + np.floating, + ) + check( + assert_type(s.shift(-1, fill_value=0), "pd.Series[int]"), pd.Series, np.integer + ) + check(assert_type(s.shift(freq="1D"), "pd.Series[int]"), pd.Series, np.integer) def test_types_rank() -> None: From 3c9dc5ba6169572b6b72c876d6cf9b02cf267a71 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Sat, 8 Feb 2025 12:17:50 +0000 Subject: [PATCH 3/6] frame too --- tests/test_frame.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index f93c59a20..1166cce07 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -519,10 +519,10 @@ def test_types_shift() -> None: df = pd.DataFrame( data={"col1": [1, 1], "col2": [3, 4]}, index=pd.date_range("2020", periods=2) ) - df.shift() - df.shift(1) - df.shift(-1) - df.shift(freq="1D") + check(assert_type(df.shift(), pd.DataFrame), pd.DataFrame) + check(assert_type(df.shift(1), pd.DataFrame), pd.DataFrame) + check(assert_type(df.shift(-1), pd.DataFrame), pd.DataFrame) + check(assert_type(df.shift(freq="1D"), pd.DataFrame), pd.DataFrame) def test_types_rank() -> None: From 03ffe84c0cf418eb77519d7a28338ef195fab822 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Mon, 10 Feb 2025 16:25:38 +0000 Subject: [PATCH 4/6] type return as Series[Any] --- pandas-stubs/core/series.pyi | 2 +- tests/test_series.py | 11 ++++------- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 7f7cceea9..7f4b2ad63 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1223,7 +1223,7 @@ class Series(IndexOpsMixin[S1], NDFrame): freq: Frequency | timedelta | None = ..., axis: AxisIndex = ..., fill_value: object | None = ..., - ) -> Series[S1]: ... + ) -> Series[Any]: ... def memory_usage(self, index: _bool = ..., deep: _bool = ...) -> int: ... def isin(self, values: Iterable | Series[S1] | dict) -> Series[_bool]: ... def between( diff --git a/tests/test_series.py b/tests/test_series.py index dbeed2260..7454e68c6 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -417,19 +417,16 @@ def test_types_sort_values_with_key() -> None: def test_types_shift() -> None: s = pd.Series([1, 2, 3], index=pd.date_range("2020", periods=3)) - # Return type "Series[int]"" not quite correct - # - https://github.com/pandas-dev/pandas-stubs/issues/1111 - # - https://github.com/pandas-dev/pandas-stubs/issues/1110 - check(assert_type(s.shift(), "pd.Series[int]"), pd.Series, np.floating) + check(assert_type(s.shift(), "pd.Series[Any]"), pd.Series, np.floating) check( - assert_type(s.shift(axis=0, periods=1), "pd.Series[int]"), + assert_type(s.shift(axis=0, periods=1), "pd.Series[Any]"), pd.Series, np.floating, ) check( - assert_type(s.shift(-1, fill_value=0), "pd.Series[int]"), pd.Series, np.integer + assert_type(s.shift(-1, fill_value=0), "pd.Series[Any]"), pd.Series, np.integer ) - check(assert_type(s.shift(freq="1D"), "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.shift(freq="1D"), "pd.Series[Any]"), pd.Series, np.integer) def test_types_rank() -> None: From bb989c12e9948731fb709f4fb7c07b9fc7eec752 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Mon, 10 Feb 2025 16:45:36 +0000 Subject: [PATCH 5/6] use just Series --- pandas-stubs/core/series.pyi | 2 +- tests/test_series.py | 8 +++----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 7f4b2ad63..393ad4ccc 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1223,7 +1223,7 @@ class Series(IndexOpsMixin[S1], NDFrame): freq: Frequency | timedelta | None = ..., axis: AxisIndex = ..., fill_value: object | None = ..., - ) -> Series[Any]: ... + ) -> Series: ... def memory_usage(self, index: _bool = ..., deep: _bool = ...) -> int: ... def isin(self, values: Iterable | Series[S1] | dict) -> Series[_bool]: ... def between( diff --git a/tests/test_series.py b/tests/test_series.py index 7454e68c6..9a1373229 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -417,16 +417,14 @@ def test_types_sort_values_with_key() -> None: def test_types_shift() -> None: s = pd.Series([1, 2, 3], index=pd.date_range("2020", periods=3)) - check(assert_type(s.shift(), "pd.Series[Any]"), pd.Series, np.floating) + check(assert_type(s.shift(), pd.Series), pd.Series, np.floating) check( assert_type(s.shift(axis=0, periods=1), "pd.Series[Any]"), pd.Series, np.floating, ) - check( - assert_type(s.shift(-1, fill_value=0), "pd.Series[Any]"), pd.Series, np.integer - ) - check(assert_type(s.shift(freq="1D"), "pd.Series[Any]"), pd.Series, np.integer) + check(assert_type(s.shift(-1, fill_value=0), pd.Series), pd.Series, np.integer) + check(assert_type(s.shift(freq="1D"), pd.Series), pd.Series, np.integer) def test_types_rank() -> None: From ced86b289fab4f4f4c5848163caade475b322c1d Mon Sep 17 00:00:00 2001 From: Marco Edward Gorelli Date: Mon, 10 Feb 2025 16:52:29 +0000 Subject: [PATCH 6/6] Update tests/test_series.py --- tests/test_series.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_series.py b/tests/test_series.py index 9a1373229..2f9b0d63b 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -419,7 +419,7 @@ def test_types_shift() -> None: s = pd.Series([1, 2, 3], index=pd.date_range("2020", periods=3)) check(assert_type(s.shift(), pd.Series), pd.Series, np.floating) check( - assert_type(s.shift(axis=0, periods=1), "pd.Series[Any]"), + assert_type(s.shift(axis=0, periods=1), pd.Series), pd.Series, np.floating, )