GH969 Fix timestamp series and offset additions

pandas-dev · Sep 14, 2024 · 7075c1e · 7075c1e
1 parent 708c8aa
commit 7075c1e
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 3 deletions.
diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi
@@ -57,7 +57,7 @@ class DatetimeIndex(DatetimeTimedeltaMixin[Timestamp], DatetimeIndexProperties):
     def __add__(self, other: TimedeltaSeries) -> TimestampSeries: ...
     @overload
     def __add__(
-        self, other: timedelta | Timedelta | TimedeltaIndex
+        self, other: timedelta | Timedelta | TimedeltaIndex | BaseOffset
     ) -> DatetimeIndex: ...
     @overload
     def __sub__(self, other: TimedeltaSeries) -> TimestampSeries: ...

diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
@@ -2073,7 +2073,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
 class TimestampSeries(Series[Timestamp]):
     @property
     def dt(self) -> TimestampProperties: ...  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
-    def __add__(self, other: TimedeltaSeries | np.timedelta64 | timedelta) -> TimestampSeries: ...  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
+    def __add__(self, other: TimedeltaSeries | np.timedelta64 | timedelta | BaseOffset) -> TimestampSeries: ...  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
     def __radd__(self, other: TimedeltaSeries | np.timedelta64 | timedelta) -> TimestampSeries: ...  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
     @overload  # type: ignore[override]
     def __sub__(
@@ -2082,7 +2082,9 @@ class TimestampSeries(Series[Timestamp]):
     @overload
     def __sub__(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
-        other: timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64,
+        other: (
+            timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64 | BaseOffset
+        ),
     ) -> TimestampSeries: ...
     def __mul__(self, other: float | Series[int] | Series[float] | Sequence[float]) -> TimestampSeries: ...  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
     def __truediv__(self, other: float | Series[int] | Series[float] | Sequence[float]) -> TimestampSeries: ...  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]

diff --git a/tests/test_series.py b/tests/test_series.py
@@ -40,6 +40,7 @@
 
 from pandas._libs.missing import NAType
 from pandas._libs.tslibs import BaseOffset
+from pandas._libs.tslibs.offsets import YearEnd
 from pandas._typing import (
     DtypeObj,
     Scalar,
@@ -3030,6 +3031,20 @@ def test_timedeltaseries_operators() -> None:
     )
 
 
+def test_timestamp_series() -> None:
+    series = pd.Series([pd.Timestamp(2024, 4, 4)])
+    check(
+        assert_type(series + YearEnd(0), TimestampSeries),
+        TimestampSeries,
+        pd.Timestamp,
+    )
+    check(
+        assert_type(series - YearEnd(0), TimestampSeries),
+        TimestampSeries,
+        pd.Timestamp,
+    )
+
+
 def test_pipe() -> None:
     ser = pd.Series(range(10))
 

diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py
@@ -750,6 +750,12 @@ def test_some_offsets() -> None:
     )
 
 
+def test_timestampseries_offset() -> None:
+    """Test that adding an offset to a timestamp series works."""
+    vv = pd.bdate_range("2024-09-01", "2024-09-10")
+    vv + pd.tseries.offsets.YearEnd(0)
+
+
 def test_types_to_numpy() -> None:
     td_s = pd.to_timedelta(pd.Series([10, 20]), "minutes")
     check(assert_type(td_s.to_numpy(), np.ndarray), np.ndarray)