Skip to content

Commit 69710a1

Browse files
bashtageKevin Sheppard
bashtage
and
Kevin Sheppard
authored
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]>
1 parent 5f86351 commit 69710a1

File tree

5 files changed

+500
-37
lines changed

5 files changed

+500
-37
lines changed

‎pandas-stubs/_libs/tslibs/period.pyi

+160-33
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,157 @@
1-
from typing import Any
1+
import datetime
2+
from typing import (
3+
Literal,
4+
Union,
5+
overload,
6+
)
7+
8+
import numpy as np
9+
from pandas import (
10+
Index,
11+
PeriodIndex,
12+
Series,
13+
Timedelta,
14+
TimedeltaIndex,
15+
)
16+
from pandas.core.series import (
17+
OffsetSeries,
18+
PeriodSeries,
19+
TimedeltaSeries,
20+
)
21+
from typing_extensions import TypeAlias
22+
23+
from pandas._libs.tslibs import NaTType
24+
from pandas._typing import npt
25+
26+
from .timestamps import Timestamp
227

328
class IncompatibleFrequency(ValueError): ...
429

5-
class Period:
30+
from pandas._libs.tslibs.offsets import BaseOffset
31+
32+
_PeriodAddSub: TypeAlias = Union[
33+
Timedelta, datetime.timedelta, np.timedelta64, np.int64, int, BaseOffset
34+
]
35+
36+
_PeriodFreqHow: TypeAlias = Literal[
37+
"S",
38+
"E",
39+
"start",
40+
"end",
41+
]
42+
43+
_PeriodToTimestampHow: TypeAlias = Union[
44+
_PeriodFreqHow,
45+
Literal[
46+
"Start",
47+
"Finish",
48+
"Begin",
49+
"End",
50+
"s",
51+
"e",
52+
"finish",
53+
"begin",
54+
],
55+
]
56+
57+
class PeriodMixin:
58+
@property
59+
def end_time(self) -> Timestamp: ...
60+
@property
61+
def start_time(self) -> Timestamp: ...
62+
63+
class Period(PeriodMixin):
664
def __init__(
765
self,
8-
value: Any = ...,
9-
freqstr: Any = ...,
10-
ordinal: Any = ...,
11-
year: Any = ...,
12-
month: int = ...,
13-
quarter: Any = ...,
14-
day: int = ...,
15-
hour: int = ...,
16-
minute: int = ...,
17-
second: int = ...,
66+
value: Period | str | None = ...,
67+
freq: str | BaseOffset | None = ...,
68+
ordinal: int | None = ...,
69+
year: int | None = ...,
70+
month: int | None = ...,
71+
quarter: int | None = ...,
72+
day: int | None = ...,
73+
hour: int | None = ...,
74+
minute: int | None = ...,
75+
second: int | None = ...,
1876
) -> None: ...
19-
def __add__(self, other) -> Period: ...
20-
def __eq__(self, other) -> bool: ...
21-
def __ge__(self, other) -> bool: ...
22-
def __gt__(self, other) -> bool: ...
23-
def __hash__(self) -> int: ...
24-
def __le__(self, other) -> bool: ...
25-
def __lt__(self, other) -> bool: ...
26-
def __new__(cls, *args, **kwargs) -> Period: ...
27-
def __ne__(self, other) -> bool: ...
28-
def __radd__(self, other) -> Period: ...
29-
def __reduce__(self, *args, **kwargs) -> Any: ... # what should this be?
30-
def __rsub__(self, other) -> Period: ...
31-
def __setstate__(self, *args, **kwargs) -> Any: ... # what should this be?
77+
@overload
78+
def __sub__(self, other: _PeriodAddSub) -> Period: ...
79+
@overload
80+
def __sub__(self, other: Period) -> BaseOffset: ...
81+
@overload
82+
def __sub__(self, other: NaTType) -> NaTType: ...
83+
@overload
84+
def __sub__(self, other: PeriodIndex) -> Index: ...
85+
@overload
86+
def __sub__(self, other: TimedeltaSeries) -> PeriodSeries: ...
87+
@overload
88+
def __sub__(self, other: TimedeltaIndex) -> PeriodIndex: ...
89+
@overload
90+
def __add__(self, other: _PeriodAddSub) -> Period: ...
91+
@overload
92+
def __add__(self, other: NaTType) -> NaTType: ...
93+
@overload
94+
def __add__(self, other: Index) -> PeriodIndex: ...
95+
@overload
96+
def __add__(self, other: OffsetSeries | TimedeltaSeries) -> PeriodSeries: ...
97+
# ignore[misc] here because we know all other comparisons
98+
# are False, so we use Literal[False]
99+
@overload
100+
def __eq__(self, other: Period) -> bool: ... # type: ignore[misc]
101+
@overload
102+
def __eq__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... # type: ignore[misc]
103+
@overload
104+
def __eq__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ... # type: ignore[misc]
105+
@overload
106+
def __eq__(self, other: object) -> Literal[False]: ...
107+
@overload
108+
def __ge__(self, other: Period) -> bool: ...
109+
@overload
110+
def __ge__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
111+
@overload
112+
def __ge__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ...
113+
@overload
114+
def __gt__(self, other: Period) -> bool: ...
115+
@overload
116+
def __gt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
117+
@overload
118+
def __gt__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ...
119+
@overload
120+
def __le__(self, other: Period) -> bool: ...
121+
@overload
122+
def __le__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
123+
@overload
124+
def __le__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ...
125+
@overload
126+
def __lt__(self, other: Period) -> bool: ...
127+
@overload
128+
def __lt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
129+
@overload
130+
def __lt__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ...
131+
# ignore[misc] here because we know all other comparisons
132+
# are False, so we use Literal[False]
133+
@overload
134+
def __ne__(self, other: Period) -> bool: ... # type: ignore[misc]
135+
@overload
136+
def __ne__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... # type: ignore[misc]
137+
@overload
138+
def __ne__(self, other: PeriodSeries | Series[Period]) -> Series[bool]: ... # type: ignore[misc]
139+
@overload
140+
def __ne__(self, other: object) -> Literal[True]: ...
141+
# Ignored due to indecipherable error from mypy:
142+
# Forward operator "__add__" is not callable [misc]
143+
@overload
144+
def __radd__(self, other: _PeriodAddSub) -> Period: ... # type: ignore[misc]
145+
# Real signature is -> PeriodIndex, but conflicts with Index.__add__
146+
# Changing Index is very hard due to Index inheritance
147+
# Signatures of "__radd__" of "Period" and "__add__" of "Index"
148+
# are unsafely overlapping
149+
@overload
150+
def __radd__(self, other: Index) -> Index: ...
151+
@overload
152+
def __radd__(self, other: TimedeltaSeries) -> PeriodSeries: ...
153+
@overload
154+
def __radd__(self, other: NaTType) -> NaTType: ...
32155
@property
33156
def day(self) -> int: ...
34157
@property
@@ -42,7 +165,7 @@ class Period:
42165
@property
43166
def end_time(self) -> Timestamp: ...
44167
@property
45-
def freq(self) -> Any: ...
168+
def freq(self) -> BaseOffset: ...
46169
@property
47170
def freqstr(self) -> str: ...
48171
@property
@@ -71,12 +194,16 @@ class Period:
71194
def weekofyear(self) -> int: ...
72195
@property
73196
def year(self) -> int: ...
74-
# Static methods
197+
@property
198+
def day_of_year(self) -> int: ...
199+
@property
200+
def day_of_week(self) -> int: ...
201+
def asfreq(self, freq: str | BaseOffset, how: _PeriodFreqHow = ...) -> Period: ...
75202
@classmethod
76-
def now(cls) -> Period: ...
77-
# Methods
78-
def asfreq(self, freq: str, how: str = ...) -> Period: ...
203+
def now(cls, freq: str | BaseOffset = ...) -> Period: ...
79204
def strftime(self, fmt: str) -> str: ...
80-
def to_timestamp(self, freq: str, how: str = ...) -> Timestamp: ...
81-
82-
from .timestamps import Timestamp
205+
def to_timestamp(
206+
self,
207+
freq: str | BaseOffset | None = ...,
208+
how: _PeriodToTimestampHow = ...,
209+
) -> Timestamp: ...

‎pandas-stubs/core/indexes/period.pyi

+15-2
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,21 @@
1-
from typing import Hashable
1+
from typing import (
2+
Hashable,
3+
overload,
4+
)
25

36
import numpy as np
47
import pandas as pd
8+
from pandas import Index
59
from pandas.core.indexes.datetimelike import (
610
DatetimeIndexOpsMixin as DatetimeIndexOpsMixin,
711
)
812
from pandas.core.indexes.numeric import Int64Index
13+
from pandas.core.series import OffsetSeries
914

10-
from pandas._libs.tslibs import BaseOffset
15+
from pandas._libs.tslibs import (
16+
BaseOffset,
17+
Period,
18+
)
1119

1220
class PeriodIndex(DatetimeIndexOpsMixin, Int64Index):
1321
def __new__(
@@ -24,6 +32,11 @@ class PeriodIndex(DatetimeIndexOpsMixin, Int64Index):
2432
@property
2533
def values(self): ...
2634
def __contains__(self, key) -> bool: ...
35+
# Override due to supertype incompatibility which has it for NumericIndex or complex.
36+
@overload # type: ignore[override]
37+
def __sub__(self, other: Period) -> Index: ...
38+
@overload
39+
def __sub__(self, other: PeriodIndex) -> OffsetSeries: ...
2740
def __array__(self, dtype=...) -> np.ndarray: ...
2841
def __array_wrap__(self, result, context=...): ...
2942
def asof_locs(self, where, mask): ...

‎pandas-stubs/core/indexes/timedeltas.pyi

+7-1
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,14 @@ from typing import (
44
overload,
55
)
66

7-
from pandas import DateOffset
7+
from pandas import (
8+
DateOffset,
9+
Period,
10+
)
811
from pandas.core.indexes.accessors import TimedeltaIndexProperties
912
from pandas.core.indexes.datetimelike import DatetimeTimedeltaMixin
1013
from pandas.core.indexes.datetimes import DatetimeIndex
14+
from pandas.core.indexes.period import PeriodIndex
1115
from pandas.core.series import TimedeltaSeries
1216

1317
from pandas._libs import (
@@ -33,6 +37,8 @@ class TimedeltaIndex(DatetimeTimedeltaMixin, TimedeltaIndexProperties):
3337
# various ignores needed for mypy, as we do want to restrict what can be used in
3438
# arithmetic for these types
3539
@overload # type: ignore[override]
40+
def __add__(self, other: Period) -> PeriodIndex: ...
41+
@overload
3642
def __add__(self, other: DatetimeIndex) -> DatetimeIndex: ...
3743
@overload
3844
def __add__(self, other: Timedelta | TimedeltaIndex) -> TimedeltaIndex: ...

‎pandas-stubs/core/series.pyi

+22-1
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,17 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
194194
name: Hashable | None = ...,
195195
copy: bool = ...,
196196
fastpath: bool = ...,
197-
) -> Series[Period]: ...
197+
) -> PeriodSeries: ...
198+
@overload
199+
def __new__(
200+
cls,
201+
data: TimedeltaIndex,
202+
index: Axes | None = ...,
203+
dtype=...,
204+
name: Hashable | None = ...,
205+
copy: bool = ...,
206+
fastpath: bool = ...,
207+
) -> TimedeltaSeries: ...
198208
@overload
199209
def __new__(
200210
cls,
@@ -1727,6 +1737,10 @@ class TimestampSeries(Series[Timestamp]):
17271737

17281738
class TimedeltaSeries(Series[Timedelta]):
17291739
# ignores needed because of mypy
1740+
@overload # type: ignore[override]
1741+
def __add__(self, other: Period) -> PeriodSeries: ...
1742+
@overload
1743+
def __add__(self, other: Timestamp | DatetimeIndex) -> TimestampSeries: ...
17301744
def __radd__(self, pther: Timestamp | TimestampSeries) -> TimestampSeries: ... # type: ignore[override]
17311745
def __mul__(self, other: num) -> TimedeltaSeries: ... # type: ignore[override]
17321746
def __sub__( # type: ignore[override]
@@ -1739,3 +1753,10 @@ class PeriodSeries(Series[Period]):
17391753
# ignore needed because of mypy
17401754
@property
17411755
def dt(self) -> PeriodProperties: ... # type: ignore[override]
1756+
def __sub__(self, other: PeriodSeries) -> OffsetSeries: ... # type: ignore[override]
1757+
1758+
class OffsetSeries(Series):
1759+
@overload # type: ignore[override]
1760+
def __radd__(self, other: Period) -> PeriodSeries: ...
1761+
@overload
1762+
def __radd__(self, other: BaseOffset) -> OffsetSeries: ...

0 commit comments

Comments
 (0)