GH1291 Add BaseOffset for shift operators on Series and DataFrame (#1301)

Author: loicdiridollou

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: ...

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: ...

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:

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: