diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index b1f734097..b0de6f5a3 100644 --- a/pandas-stubs/core/indexes/datetimes.pyi +++ b/pandas-stubs/core/indexes/datetimes.pyi @@ -57,13 +57,13 @@ 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: ... @overload def __sub__( - self, other: timedelta | Timedelta | TimedeltaIndex + self, other: timedelta | Timedelta | TimedeltaIndex | BaseOffset ) -> DatetimeIndex: ... @overload def __sub__( diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 24bac1a11..24c475762 100644 --- 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 index 5a2110960..210e4bc5d 100644 --- 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 index eabbc3707..1d7a11bf7 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -750,6 +750,13 @@ 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") + shifted_vv = vv + pd.tseries.offsets.YearEnd(0) + check(assert_type(shifted_vv, pd.DatetimeIndex), pd.DatetimeIndex) + + 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)