Attempt to introduce OffsetSeries

Author: bashtage

pandas-stubs/_libs/tslibs/period.pyi

@@ -13,6 +13,7 @@ from pandas import (
     TimedeltaIndex,
 )
 from pandas.core.series import (
+    OffsetSeries,
     PeriodSeries,
     TimedeltaSeries,
 )
@@ -91,7 +92,7 @@ class Period(PeriodMixin):
     @overload
     def __add__(self, other: Index) -> PeriodIndex: ...
     @overload
-    def __add__(self, other: TimedeltaSeries) -> PeriodSeries: ...
+    def __add__(self, other: OffsetSeries | TimedeltaSeries) -> PeriodSeries: ...
     @overload  # type: ignore[override]
     def __eq__(self, other: Period) -> bool: ...
     @overload

pandas-stubs/core/indexes/period.pyi

@@ -1,4 +1,7 @@
-from typing import Hashable
+from typing import (
+    Hashable,
+    overload,
+)
 
 import numpy as np
 import pandas as pd
@@ -7,6 +10,7 @@ from pandas.core.indexes.datetimelike import (
     DatetimeIndexOpsMixin as DatetimeIndexOpsMixin,
 )
 from pandas.core.indexes.numeric import Int64Index
+from pandas.core.series import OffsetSeries
 
 from pandas._libs.tslibs import (
     BaseOffset,
@@ -28,8 +32,11 @@ class PeriodIndex(DatetimeIndexOpsMixin, Int64Index):
     @property
     def values(self): ...
     def __contains__(self, key) -> bool: ...
-    # Override due to supertpye incompatibility which has it for NumericIndex or complex.
-    def __sub__(self, other: Period) -> Index: ...  # type: ignore[override]
+    # Override due to supertype incompatibility which has it for NumericIndex or complex.
+    @overload  # type: ignore[override]
+    def __sub__(self, other: Period) -> Index: ...
+    @overload
+    def __sub__(self, other: PeriodIndex) -> OffsetSeries: ...
     def __array__(self, dtype=...) -> np.ndarray: ...
     def __array_wrap__(self, result, context=...): ...
     def asof_locs(self, where, mask): ...

pandas-stubs/core/series.pyi

@@ -194,7 +194,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
         name: Hashable | None = ...,
         copy: bool = ...,
         fastpath: bool = ...,
-    ) -> Series[Period]: ...
+    ) -> PeriodSeries: ...
     @overload
     def __new__(
         cls,
@@ -1749,3 +1749,10 @@ class PeriodSeries(Series[Period]):
     # ignore needed because of mypy
     @property
     def dt(self) -> PeriodProperties: ...  # type: ignore[override]
+    def __sub__(self, other: PeriodSeries) -> OffsetSeries: ...  # type: ignore[override]
+
+class OffsetSeries(Series):
+    @overload  # type: ignore[override]
+    def __radd__(self, other: Period) -> PeriodSeries: ...
+    @overload
+    def __radd__(self, other: BaseOffset) -> OffsetSeries: ...

tests/test_scalars.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
 import datetime as dt
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    cast,
-)
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
-from typing_extensions import assert_type
+from typing_extensions import (
+    TypeAlias,
+    assert_type,
+)
 
 from pandas._libs.tslibs import (
     BaseOffset,
@@ -18,13 +17,20 @@
 
 if TYPE_CHECKING:
     from pandas.core.series import (
+        OffsetSeries,
         PeriodSeries,
         TimedeltaSeries,
     )
 
     from pandas._typing import np_ndarray_bool
 else:
-    PeriodSeries = TimedeltaSeries = np_ndarray_bool = Any
+    import numpy.typing as npt
+
+    np_ndarray_bool = npt.NDArray[np.bool_]
+    PeriodSeries: TypeAlias = pd.Series
+    TimedeltaSeries: TypeAlias = pd.Series
+    OffsetSeries: TypeAlias = pd.Series
+
 
 from tests import check
 
@@ -100,11 +106,13 @@ def test_periof_add_subtract() -> None:
     as_np_i64 = np.int64(1)
     as_int = int(1)
     as_period_index = pd.period_range("2012-1-1", periods=10, freq="D")
+    check(assert_type(as_period_index, pd.PeriodIndex), pd.PeriodIndex)
     as_period = pd.Period("2012-1-1", freq="D")
     scale = 24 * 60 * 60 * 10**9
     as_td_series = pd.Series(pd.timedelta_range(scale, scale, freq="D"))
     check(assert_type(as_td_series, TimedeltaSeries), pd.Series)
     as_period_series = pd.Series(as_period_index)
+    check(assert_type(as_period_series, PeriodSeries), pd.Series)
     as_timedelta_idx = pd.timedelta_range(scale, scale, freq="D")
     as_nat = pd.NaT
 
@@ -114,18 +122,21 @@ def test_periof_add_subtract() -> None:
     check(assert_type(p + as_np_i64, pd.Period), pd.Period)
     check(assert_type(p + as_int, pd.Period), pd.Period)
     check(assert_type(p + p.freq, pd.Period), pd.Period)
-    check(assert_type(p + (p - as_period_index), pd.PeriodIndex), pd.PeriodIndex)
+    # offset_index is tested below
+    offset_index = p - as_period_index
+    check(assert_type(p + offset_index, pd.PeriodIndex), pd.PeriodIndex)
     check(assert_type(p + as_td_series, PeriodSeries), pd.Series, pd.Period)
     check(assert_type(p + as_timedelta_idx, pd.PeriodIndex), pd.PeriodIndex)
     check(assert_type(p + as_nat, NaTType), NaTType)
-    das8 = cast(TimedeltaSeries, (as_period_series - as_period_series))
-    check(assert_type(p + das8, PeriodSeries), pd.Series, pd.Period)
+    offset_series = as_period_series - as_period_series
+    check(assert_type(offset_series, OffsetSeries), pd.Series)
+    check(assert_type(p + offset_series, PeriodSeries), pd.Series, pd.Period)
     check(assert_type(p - as_pd_td, pd.Period), pd.Period)
     check(assert_type(p - as_dt_td, pd.Period), pd.Period)
     check(assert_type(p - as_np_td, pd.Period), pd.Period)
     check(assert_type(p - as_np_i64, pd.Period), pd.Period)
     check(assert_type(p - as_int, pd.Period), pd.Period)
-    check(assert_type(p - as_period_index, pd.Index), pd.Index)
+    check(assert_type(offset_index, pd.Index), pd.Index)
     check(assert_type(p - as_period, BaseOffset), Day)
     check(assert_type(p - as_td_series, PeriodSeries), pd.Series, pd.Period)
     check(assert_type(p - as_timedelta_idx, pd.PeriodIndex), pd.PeriodIndex)