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: