ENH: Improve typing for Period (#390)

* ENH: Improve typing for Period

* ENH: Improve list of how values

* ENH: Improve Period typing and testing

* ENH: Improve dtype typing and testing

* REV: Revert unrelated changes

* Incorporate feedback for PR

* Attempt to introduce OffsetSeries

* ENH: Futher improvements to Period

* ENH: Add support for Series

Co-authored-by: Kevin Sheppard <[email protected]>

pandas-stubs/_libs/tslibs/period.pyi

@@ -1,34 +1,157 @@
-from typing import Any
+import datetime
+from typing import (
+    Literal,
+    Union,
+    overload,
+)
+
+import numpy as np
+from pandas import (
+    Index,
+    PeriodIndex,
+    Series,
+    Timedelta,
+    TimedeltaIndex,
+)
+from pandas.core.series import (
+    OffsetSeries,
+    PeriodSeries,
+    TimedeltaSeries,
+)
+from typing_extensions import TypeAlias
+
+from pandas._libs.tslibs import NaTType
+from pandas._typing import npt
+
+from .timestamps import Timestamp
 
 class IncompatibleFrequency(ValueError): ...
 
-class Period:
+from pandas._libs.tslibs.offsets import BaseOffset
+
+_PeriodAddSub: TypeAlias = Union[
+    Timedelta, datetime.timedelta, np.timedelta64, np.int64, int, BaseOffset
+]
+
+_PeriodFreqHow: TypeAlias = Literal[
+    "S",
+    "E",
+    "start",
+    "end",
+]
+
+_PeriodToTimestampHow: TypeAlias = Union[
+    _PeriodFreqHow,
+    Literal[
+        "Start",
+        "Finish",
+        "Begin",
+        "End",
+        "s",
+        "e",
+        "finish",
+        "begin",
+    ],
+]
+
+class PeriodMixin:
+    @property
+    def end_time(self) -> Timestamp: ...
+    @property
+    def start_time(self) -> Timestamp: ...
+
+class Period(PeriodMixin):
     def __init__(
         self,
-        value: Any = ...,
-        freqstr: Any = ...,
-        ordinal: Any = ...,
-        year: Any = ...,
-        month: int = ...,
-        quarter: Any = ...,
-        day: int = ...,
-        hour: int = ...,
-        minute: int = ...,
-        second: int = ...,
+        value: Period | str | None = ...,
+        freq: str | BaseOffset | None = ...,
+        ordinal: int | None = ...,
+        year: int | None = ...,
+        month: int | None = ...,
+        quarter: int | None = ...,
+        day: int | None = ...,
+        hour: int | None = ...,
+        minute: int | None = ...,
+        second: int | None = ...,
     ) -> None: ...
-    def __add__(self, other) -> Period: ...
-    def __eq__(self, other) -> bool: ...
-    def __ge__(self, other) -> bool: ...
-    def __gt__(self, other) -> bool: ...
-    def __hash__(self) -> int: ...
-    def __le__(self, other) -> bool: ...
-    def __lt__(self, other) -> bool: ...
-    def __new__(cls, *args, **kwargs) -> Period: ...
-    def __ne__(self, other) -> bool: ...
-    def __radd__(self, other) -> Period: ...
-    def __reduce__(self, *args, **kwargs) -> Any: ...  # what should this be?
-    def __rsub__(self, other) -> Period: ...
-    def __setstate__(self, *args, **kwargs) -> Any: ...  # what should this be?
+    @overload
+    def __sub__(self, other: _PeriodAddSub) -> Period: ...
+    @overload
+    def __sub__(self, other: Period) -> BaseOffset: ...
+    @overload
+    def __sub__(self, other: NaTType) -> NaTType: ...
+    @overload
+    def __sub__(self, other: PeriodIndex) -> Index: ...
+    @overload
+    def __sub__(self, other: TimedeltaSeries) -> PeriodSeries: ...
+    @overload
+    def __sub__(self, other: TimedeltaIndex) -> PeriodIndex: ...
+    @overload
+    def __add__(self, other: _PeriodAddSub) -> Period: ...
+    @overload
+    def __add__(self, other: NaTType) -> NaTType: ...
+    @overload
+    def __add__(self, other: Index) -> PeriodIndex: ...
+    @overload
+    def __add__(self, other: OffsetSeries | TimedeltaSeries) -> PeriodSeries: ...
+    #  ignore[misc] here because we know all other comparisons
+    #  are False, so we use Literal[False]
+    @overload
+    def __eq__(self, other: Period) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __eq__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...  # type: ignore[misc]
+    @overload
+    def __eq__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ...  # type: ignore[misc]
+    @overload
+    def __eq__(self, other: object) -> Literal[False]: ...
+    @overload
+    def __ge__(self, other: Period) -> bool: ...
+    @overload
+    def __ge__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
+    @overload
+    def __ge__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ...
+    @overload
+    def __gt__(self, other: Period) -> bool: ...
+    @overload
+    def __gt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
+    @overload
+    def __gt__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ...
+    @overload
+    def __le__(self, other: Period) -> bool: ...
+    @overload
+    def __le__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
+    @overload
+    def __le__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ...
+    @overload
+    def __lt__(self, other: Period) -> bool: ...
+    @overload
+    def __lt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
+    @overload
+    def __lt__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ...
+    #  ignore[misc] here because we know all other comparisons
+    #  are False, so we use Literal[False]
+    @overload
+    def __ne__(self, other: Period) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __ne__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...  # type: ignore[misc]
+    @overload
+    def __ne__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ...  # type: ignore[misc]
+    @overload
+    def __ne__(self, other: object) -> Literal[True]: ...
+    # Ignored due to indecipherable error from mypy:
+    # Forward operator "__add__" is not callable  [misc]
+    @overload
+    def __radd__(self, other: _PeriodAddSub) -> Period: ...  # type: ignore[misc]
+    # Real signature is -> PeriodIndex, but conflicts with Index.__add__
+    # Changing Index is very hard due to Index inheritance
+    #   Signatures of "__radd__" of "Period" and "__add__" of "Index"
+    #   are unsafely overlapping
+    @overload
+    def __radd__(self, other: Index) -> Index: ...
+    @overload
+    def __radd__(self, other: TimedeltaSeries) -> PeriodSeries: ...
+    @overload
+    def __radd__(self, other: NaTType) -> NaTType: ...
     @property
     def day(self) -> int: ...
     @property
@@ -42,7 +165,7 @@ class Period:
     @property
     def end_time(self) -> Timestamp: ...
     @property
-    def freq(self) -> Any: ...
+    def freq(self) -> BaseOffset: ...
     @property
     def freqstr(self) -> str: ...
     @property
@@ -71,12 +194,16 @@ class Period:
     def weekofyear(self) -> int: ...
     @property
     def year(self) -> int: ...
-    # Static methods
+    @property
+    def day_of_year(self) -> int: ...
+    @property
+    def day_of_week(self) -> int: ...
+    def asfreq(self, freq: str | BaseOffset, how: _PeriodFreqHow = ...) -> Period: ...
     @classmethod
-    def now(cls) -> Period: ...
-    # Methods
-    def asfreq(self, freq: str, how: str = ...) -> Period: ...
+    def now(cls, freq: str | BaseOffset = ...) -> Period: ...
     def strftime(self, fmt: str) -> str: ...
-    def to_timestamp(self, freq: str, how: str = ...) -> Timestamp: ...
-
-from .timestamps import Timestamp
+    def to_timestamp(
+        self,
+        freq: str | BaseOffset | None = ...,
+        how: _PeriodToTimestampHow = ...,
+    ) -> Timestamp: ...

pandas-stubs/core/indexes/period.pyi

@@ -1,13 +1,21 @@
-from typing import Hashable
+from typing import (
+    Hashable,
+    overload,
+)
 
 import numpy as np
 import pandas as pd
+from pandas import Index
 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
+from pandas._libs.tslibs import (
+    BaseOffset,
+    Period,
+)
 
 class PeriodIndex(DatetimeIndexOpsMixin, Int64Index):
     def __new__(
@@ -24,6 +32,11 @@ class PeriodIndex(DatetimeIndexOpsMixin, Int64Index):
     @property
     def values(self): ...
     def __contains__(self, key) -> bool: ...
+    # 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/indexes/timedeltas.pyi

@@ -4,10 +4,14 @@ from typing import (
     overload,
 )
 
-from pandas import DateOffset
+from pandas import (
+    DateOffset,
+    Period,
+)
 from pandas.core.indexes.accessors import TimedeltaIndexProperties
 from pandas.core.indexes.datetimelike import DatetimeTimedeltaMixin
 from pandas.core.indexes.datetimes import DatetimeIndex
+from pandas.core.indexes.period import PeriodIndex
 from pandas.core.series import TimedeltaSeries
 
 from pandas._libs import (
@@ -33,6 +37,8 @@ class TimedeltaIndex(DatetimeTimedeltaMixin, TimedeltaIndexProperties):
     # various ignores needed for mypy, as we do want to restrict what can be used in
     # arithmetic for these types
     @overload  # type: ignore[override]
+    def __add__(self, other: Period) -> PeriodIndex: ...
+    @overload
     def __add__(self, other: DatetimeIndex) -> DatetimeIndex: ...
     @overload
     def __add__(self, other: Timedelta | TimedeltaIndex) -> TimedeltaIndex: ...

pandas-stubs/core/series.pyi

@@ -194,7 +194,17 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
         name: Hashable | None = ...,
         copy: bool = ...,
         fastpath: bool = ...,
-    ) -> Series[Period]: ...
+    ) -> PeriodSeries: ...
+    @overload
+    def __new__(
+        cls,
+        data: TimedeltaIndex,
+        index: Axes | None = ...,
+        dtype=...,
+        name: Hashable | None = ...,
+        copy: bool = ...,
+        fastpath: bool = ...,
+    ) -> TimedeltaSeries: ...
     @overload
     def __new__(
         cls,
@@ -1727,6 +1737,10 @@ class TimestampSeries(Series[Timestamp]):
 
 class TimedeltaSeries(Series[Timedelta]):
     # ignores needed because of mypy
+    @overload  # type: ignore[override]
+    def __add__(self, other: Period) -> PeriodSeries: ...
+    @overload
+    def __add__(self, other: Timestamp | DatetimeIndex) -> TimestampSeries: ...
     def __radd__(self, pther: Timestamp | TimestampSeries) -> TimestampSeries: ...  # type: ignore[override]
     def __mul__(self, other: num) -> TimedeltaSeries: ...  # type: ignore[override]
     def __sub__(  # type: ignore[override]
@@ -1739,3 +1753,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: ...