diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 3cb74eee0..80521be87 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -892,7 +892,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): def shift( self, periods: int | Sequence[int] = ..., - freq: DateOffset | dt.timedelta | _str | None = ..., + freq: BaseOffset | dt.timedelta | _str | None = ..., axis: Axis = ..., fill_value: Scalar | NAType | None = ..., ) -> Self: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index f0dc7e7fc..a0d4e5b49 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1136,7 +1136,7 @@ class Series(IndexOpsMixin[S1], NDFrame): def shift( self, periods: int | Sequence[int] = ..., - freq: DateOffset | timedelta | _str | None = ..., + freq: BaseOffset | timedelta | _str | None = ..., axis: Axis = ..., fill_value: Scalar | NAType | None = ..., ) -> Series: ... diff --git a/tests/series/test_series.py b/tests/series/test_series.py index 49fd80962..34f888e6c 100644 --- a/tests/series/test_series.py +++ b/tests/series/test_series.py @@ -63,6 +63,10 @@ from pandas.io.formats.format import EngFormatter from pandas.tseries.offsets import ( BaseOffset, + BDay, + BQuarterEnd, + MonthEnd, + Week, YearEnd, ) @@ -448,6 +452,7 @@ def test_types_sort_values_with_key() -> None: def test_types_shift() -> None: + """Test shift operator on series with different arguments.""" 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( @@ -457,6 +462,11 @@ def test_types_shift() -> None: ) 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) + check(assert_type(s.shift(freq=BDay(1)), pd.Series), pd.Series, np.integer) + check(assert_type(s.shift(freq=BQuarterEnd(5)), pd.Series), pd.Series, np.integer) + check(assert_type(s.shift(freq=MonthEnd(3)), pd.Series), pd.Series, np.integer) + check(assert_type(s.shift(freq=Week(4)), pd.Series), pd.Series, np.integer) + check(assert_type(s.shift(freq=YearEnd(2)), pd.Series), pd.Series, np.integer) def test_series_pct_change() -> None: diff --git a/tests/test_frame.py b/tests/test_frame.py index b2348fc30..16c651c6b 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -60,6 +60,13 @@ from pandas.io.formats.format import EngFormatter from pandas.io.formats.style import Styler from pandas.io.parsers import TextFileReader +from pandas.tseries.offsets import ( + BDay, + BQuarterEnd, + MonthEnd, + Week, + YearEnd, +) if TYPE_CHECKING: from pandas.core.frame import _PandasNamedTuple @@ -603,6 +610,11 @@ def test_types_shift() -> None: 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) + check(assert_type(df.shift(freq=BDay(1)), pd.DataFrame), pd.DataFrame) + check(assert_type(df.shift(freq=BQuarterEnd(5)), pd.DataFrame), pd.DataFrame) + check(assert_type(df.shift(freq=MonthEnd(3)), pd.DataFrame), pd.DataFrame) + check(assert_type(df.shift(freq=Week(4)), pd.DataFrame), pd.DataFrame) + check(assert_type(df.shift(freq=YearEnd(2)), pd.DataFrame), pd.DataFrame) def test_types_rank() -> None: