Repository: pandas-dev/pandas-stubs (language: python)

Issue:
Check incorrect when using a TimestampSeries and YearEnd from pandas.tseries.offsets
It seems type annotations are wrong between a Timestamp series and tools offered by the offsets library, using `pyright`

df["year_end"] = df["date"] + pd.tseries.offsets.YearEnd(0) is flagged with Operator "+" not supported for types "TimestampSeries" and "YearEnd"

 - OS: MacOS
 - OS Version 13.5
 - python 3.9.19
 - pylance version v2024.7.1


Discussion / hints:
Can you provide a more complete example?
I digged into that one as I remember having similar issues on a project but I was not able to fully reproduce with `TimestampSeries`.
Instead here is a short example that can reproduce a similar error, @Dr-Irv let me know if this is something that we thrive to fix through an overload (I assume, haven't looked at the code in details).
```python
import pandas as pd

df = pd.DataFrame({"date": pd.bdate_range("2024-09-01", "2024-09-10")})
df["date"] = pd.to_datetime(df["date"])
vv = pd.bdate_range("2024-09-01", "2024-09-10")

vv + pd.tseries.offsets.YearEnd(0)
```
which will raise the error:
`Operator "+" not supported for types "DatetimeIndex" and "YearEnd"`
https://github.com/pandas-dev/pandas-stubs/blob/5177d67a3189aee4037031ffd503f0af997c617b/pandas-stubs/core/indexes/datetimes.pyi#L58-L61
Likely could be added here as `pd.tseries.offsets.YearEnd` is a class from the `BaseOffset` base class and not `Timedelta`.
Open to discussing if this is the correct approach.
Might be the same / similar as https://github.com/pandas-dev/pandas-stubs/issues/755
> I digged into that one as I remember having similar issues on a project but I was not able to fully reproduce with `TimestampSeries`. Instead here is a short example that can reproduce a similar error, @Dr-Irv let me know if this is something that we thrive to fix through an overload (I assume, haven't looked at the code in details).
> 
> ```python
> import pandas as pd
> 
> df = pd.DataFrame({"date": pd.bdate_range("2024-09-01", "2024-09-10")})
> df["date"] = pd.to_datetime(df["date"])
> vv = pd.bdate_range("2024-09-01", "2024-09-10")
> 
> vv + pd.tseries.offsets.YearEnd(0)
> ```
> 
> which will raise the error: `Operator "+" not supported for types "DatetimeIndex" and "YearEnd"`

Feel free to open a MR - I think this should work!

Fix patch:
diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi
index b1f73409..b0de6f5a 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 24bac1a1..24c47576 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]


Test patch:
diff --git a/tests/test_series.py b/tests/test_series.py
index 5a211096..210e4bc5 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 eabbc370..1d7a11bf 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)