Skip to content

Commit d876a71

Browse files
bashtagebashtage
committed
ENH: Improve typing for Period
1 parent 5f86351 commit d876a71

File tree

2 files changed

+311
-32
lines changed

2 files changed

+311
-32
lines changed

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

Lines changed: 122 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,120 @@
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+
Timedelta,
13+
)
14+
from pandas.core.series import (
15+
PeriodSeries,
16+
TimedeltaSeries,
17+
)
18+
from typing_extensions import TypeAlias
19+
20+
from pandas._typing import npt
21+
22+
from .timestamps import Timestamp
223

324
class IncompatibleFrequency(ValueError): ...
425

5-
class Period:
26+
from pandas._libs.tslibs.offsets import BaseOffset
27+
28+
_PeriodAddSub: TypeAlias = Union[
29+
Timedelta, datetime.timedelta, np.timedelta64, np.int64, int, BaseOffset
30+
]
31+
32+
_PeriodFreqHow: TypeAlias = Literal[
33+
"S",
34+
"E",
35+
"Start",
36+
"Finish",
37+
"Begin",
38+
"End",
39+
"s",
40+
"e",
41+
"start",
42+
"finish",
43+
"begin",
44+
"end",
45+
]
46+
47+
class PeriodMixin:
48+
@property
49+
def end_time(self) -> Timestamp: ...
50+
@property
51+
def start_time(self) -> Timestamp: ...
52+
53+
class Period(PeriodMixin):
654
def __init__(
755
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 = ...,
56+
value: Period | str | None = ...,
57+
freq: str | BaseOffset | None = ...,
58+
ordinal: int | None = ...,
59+
year: int | None = ...,
60+
month: int | None = ...,
61+
quarter: int | None = ...,
62+
day: int | None = ...,
63+
hour: int | None = ...,
64+
minute: int | None = ...,
65+
second: int | None = ...,
1866
) -> 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: ...
67+
@overload
68+
def __sub__(self, other: _PeriodAddSub) -> Period: ...
69+
@overload
70+
def __sub__(self, other: Period) -> BaseOffset: ...
71+
@overload
72+
def __sub__(self, other: PeriodIndex) -> Index: ...
73+
@overload
74+
def __sub__(self, other: TimedeltaSeries) -> PeriodSeries: ...
75+
@overload
76+
def __add__(self, other: _PeriodAddSub) -> Period: ...
77+
@overload
78+
def __add__(self, other: Index) -> PeriodIndex: ...
79+
@overload
80+
def __add__(self, other: TimedeltaSeries) -> PeriodSeries: ...
81+
@overload # type: ignore[override]
82+
def __eq__(self, other: Period) -> bool: ...
83+
@overload
84+
def __eq__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
85+
@overload
86+
def __ge__(self, other: Period) -> bool: ...
87+
@overload
88+
def __ge__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
89+
@overload
90+
def __gt__(self, other: Period) -> bool: ...
91+
@overload
92+
def __gt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
2393
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?
94+
@overload
95+
def __le__(self, other: Period) -> bool: ...
96+
@overload
97+
def __le__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
98+
@overload
99+
def __lt__(self, other: Period) -> bool: ...
100+
@overload
101+
def __lt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
102+
@overload # type: ignore[override]
103+
def __ne__(self, other: Period) -> bool: ...
104+
@overload
105+
def __ne__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
106+
# Ignored due to indecipherable error from mypy:
107+
# Forward operator "__add__" is not callable [misc]
108+
@overload
109+
def __radd__(self, other: _PeriodAddSub) -> Period: ... # type: ignore[misc]
110+
# Real signature is -> PeriodIndex, but conflicts with Index.__add__
111+
# Changing Index is very hard due to Index inheritance
112+
# Signatures of "__radd__" of "Period" and "__add__" of "Index"
113+
# are unsafely overlapping
114+
@overload
115+
def __radd__(self, other: Index) -> Index: ...
116+
@overload
117+
def __radd__(self, other: TimedeltaSeries) -> PeriodSeries: ...
32118
@property
33119
def day(self) -> int: ...
34120
@property
@@ -42,7 +128,7 @@ class Period:
42128
@property
43129
def end_time(self) -> Timestamp: ...
44130
@property
45-
def freq(self) -> Any: ...
131+
def freq(self) -> BaseOffset: ...
46132
@property
47133
def freqstr(self) -> str: ...
48134
@property
@@ -71,12 +157,16 @@ class Period:
71157
def weekofyear(self) -> int: ...
72158
@property
73159
def year(self) -> int: ...
74-
# Static methods
160+
@property
161+
def day_of_year(self) -> int: ...
162+
@property
163+
def day_of_week(self) -> int: ...
164+
def asfreq(self, freq: str | BaseOffset, how: _PeriodFreqHow = ...) -> Period: ...
75165
@classmethod
76-
def now(cls) -> Period: ...
77-
# Methods
78-
def asfreq(self, freq: str, how: str = ...) -> Period: ...
166+
def now(cls, freq: str | BaseOffset = ...) -> Period: ...
79167
def strftime(self, fmt: str) -> str: ...
80-
def to_timestamp(self, freq: str, how: str = ...) -> Timestamp: ...
81-
82-
from .timestamps import Timestamp
168+
def to_timestamp(
169+
self,
170+
freq: str | BaseOffset | None = ...,
171+
how: _PeriodFreqHow = ...,
172+
) -> Timestamp: ...

‎tests/test_scalars.py

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
from __future__ import annotations
2+
3+
import datetime as dt
4+
from typing import (
5+
TYPE_CHECKING,
6+
Any,
7+
cast,
8+
)
9+
10+
import numpy as np
11+
import pandas as pd
12+
from typing_extensions import assert_type
13+
14+
from pandas._libs.tslibs import BaseOffset
15+
16+
if TYPE_CHECKING:
17+
from pandas.core.series import PeriodSeries # noqa: F401
18+
from pandas.core.series import TimedeltaSeries # noqa: F401
19+
20+
from pandas._typing import np_ndarray_bool
21+
else:
22+
np_ndarray_bool = Any
23+
24+
from tests import check
25+
26+
from pandas.tseries.offsets import Day
27+
28+
29+
def test_period() -> None:
30+
p = pd.Period("2012-1-1", freq="D")
31+
check(assert_type(p, pd.Period), pd.Period)
32+
check(assert_type(pd.Period(p), pd.Period), pd.Period)
33+
check(assert_type(pd.Period("2012-1-1", freq=Day()), pd.Period), pd.Period)
34+
check(
35+
assert_type(pd.Period(freq="D", year=2012, day=1, month=1), pd.Period),
36+
pd.Period,
37+
)
38+
check(
39+
assert_type(pd.Period(None, "D", year=2012, day=1, month=1), pd.Period),
40+
pd.Period,
41+
)
42+
check(
43+
assert_type(pd.Period(None, "D", 1, year=2012, day=1, month=1), pd.Period),
44+
pd.Period,
45+
)
46+
check(
47+
assert_type(
48+
pd.Period(
49+
freq="s", year=2012, month=1, day=1, hour=12, minute=30, second=45
50+
),
51+
pd.Period,
52+
),
53+
pd.Period,
54+
)
55+
check(assert_type(pd.Period(freq="Q", year=2012, quarter=2), pd.Period), pd.Period)
56+
check(assert_type(p.day, int), int)
57+
check(assert_type(p.day_of_week, int), int)
58+
check(assert_type(p.day_of_year, int), int)
59+
check(assert_type(p.dayofweek, int), int)
60+
check(assert_type(p.dayofyear, int), int)
61+
check(assert_type(p.days_in_month, int), int)
62+
check(assert_type(p.daysinmonth, int), int)
63+
check(assert_type(p.end_time, pd.Timestamp), pd.Timestamp)
64+
check(assert_type(p.freqstr, str), str)
65+
check(assert_type(p.hour, int), int)
66+
check(assert_type(p.is_leap_year, bool), bool)
67+
check(assert_type(p.minute, int), int)
68+
check(assert_type(p.month, int), int)
69+
check(assert_type(p.quarter, int), int)
70+
check(assert_type(p.qyear, int), int)
71+
check(assert_type(p.second, int), int)
72+
check(assert_type(p.start_time, pd.Timestamp), pd.Timestamp)
73+
check(assert_type(p.week, int), int)
74+
check(assert_type(p.weekday, int), int)
75+
check(assert_type(p.weekofyear, int), int)
76+
check(assert_type(p.year, int), int)
77+
check(assert_type(p.freq, BaseOffset), Day)
78+
check(assert_type(p.ordinal, int), int)
79+
80+
p2 = pd.Period("2012-1-1", freq="2D")
81+
check(assert_type(p2.freq, BaseOffset), Day)
82+
83+
as0 = pd.Timedelta(1, "D")
84+
as1 = dt.timedelta(days=1)
85+
as2 = np.timedelta64(1, "D")
86+
as3 = np.int64(1)
87+
as4 = int(1)
88+
as5 = pd.period_range("2012-1-1", periods=10, freq="D")
89+
as6 = pd.Period("2012-1-1", freq="D")
90+
as7 = cast("TimedeltaSeries", pd.Series([pd.Timedelta(days=1)]))
91+
as8 = cast("PeriodSeries", pd.Series([as6]))
92+
93+
check(assert_type(p + as0, pd.Period), pd.Period)
94+
check(assert_type(p + as1, pd.Period), pd.Period)
95+
check(assert_type(p + as2, pd.Period), pd.Period)
96+
check(assert_type(p + as3, pd.Period), pd.Period)
97+
check(assert_type(p + as4, pd.Period), pd.Period)
98+
check(assert_type(p + p.freq, pd.Period), pd.Period)
99+
check(assert_type(p + (p - as5), pd.PeriodIndex), pd.PeriodIndex)
100+
check(assert_type(p + as7, "PeriodSeries"), pd.Series)
101+
das8 = cast("TimedeltaSeries", (as8 - as8))
102+
check(assert_type(p + das8, "PeriodSeries"), pd.Series)
103+
check(assert_type(p - as0, pd.Period), pd.Period)
104+
check(assert_type(p - as1, pd.Period), pd.Period)
105+
check(assert_type(p - as2, pd.Period), pd.Period)
106+
check(assert_type(p - as3, pd.Period), pd.Period)
107+
check(assert_type(p - as4, pd.Period), pd.Period)
108+
check(assert_type(p - as5, pd.Index), pd.Index)
109+
check(assert_type(p - as6, BaseOffset), Day)
110+
check(assert_type(p - as7, "PeriodSeries"), pd.Series)
111+
check(assert_type(p - p.freq, pd.Period), pd.Period)
112+
113+
check(assert_type(as0 + p, pd.Period), pd.Period)
114+
check(assert_type(as1 + p, pd.Period), pd.Period)
115+
check(assert_type(as2 + p, pd.Period), pd.Period)
116+
check(assert_type(as3 + p, pd.Period), pd.Period)
117+
check(assert_type(as4 + p, pd.Period), pd.Period)
118+
check(assert_type(as7 + p, "PeriodSeries"), pd.Series)
119+
check(assert_type(p.freq + p, pd.Period), pd.Period)
120+
# TODO: PeriodIndex should have a __sub__ with correct types, this op is valid
121+
# and so the assert_type is skipped
122+
check(as5 - p, pd.Index) # type: ignore[operator]
123+
124+
check(assert_type(p.__radd__(as0), pd.Period), pd.Period)
125+
check(assert_type(p.__radd__(as1), pd.Period), pd.Period)
126+
check(assert_type(p.__radd__(as2), pd.Period), pd.Period)
127+
check(assert_type(p.__radd__(as3), pd.Period), pd.Period)
128+
check(assert_type(p.__radd__(as4), pd.Period), pd.Period)
129+
check(assert_type(p.__radd__(p.freq), pd.Period), pd.Period)
130+
131+
c0 = pd.Period("2012-1-1", freq="D")
132+
c1 = pd.period_range("2012-1-1", periods=10, freq="D")
133+
134+
check(assert_type(p == c0, bool), bool)
135+
check(assert_type(p == c1, np_ndarray_bool), np.ndarray)
136+
check(assert_type(c0 == p, bool), bool)
137+
check(assert_type(c1 == p, np_ndarray_bool), np.ndarray)
138+
139+
check(assert_type(p != c0, bool), bool)
140+
check(assert_type(p != c1, np_ndarray_bool), np.ndarray)
141+
check(assert_type(c0 != p, bool), bool)
142+
check(assert_type(c1 != p, np_ndarray_bool), np.ndarray)
143+
144+
check(assert_type(p > c0, bool), bool)
145+
check(assert_type(p > c1, np_ndarray_bool), np.ndarray)
146+
check(assert_type(c0 > p, bool), bool)
147+
check(assert_type(c1 > p, np_ndarray_bool), np.ndarray)
148+
149+
check(assert_type(p < c0, bool), bool)
150+
check(assert_type(p < c1, np_ndarray_bool), np.ndarray)
151+
check(assert_type(c0 < p, bool), bool)
152+
check(assert_type(c1 < p, np_ndarray_bool), np.ndarray)
153+
154+
check(assert_type(p <= c0, bool), bool)
155+
check(assert_type(p <= c1, np_ndarray_bool), np.ndarray)
156+
check(assert_type(c0 <= p, bool), bool)
157+
check(assert_type(c1 <= p, np_ndarray_bool), np.ndarray)
158+
159+
check(assert_type(p >= c0, bool), bool)
160+
check(assert_type(p >= c1, np_ndarray_bool), np.ndarray)
161+
check(assert_type(c0 >= p, bool), bool)
162+
check(assert_type(c1 >= p, np_ndarray_bool), np.ndarray)
163+
164+
p3 = pd.Period("2007-01", freq="M")
165+
check(assert_type(p3.to_timestamp("D", "S"), pd.Timestamp), pd.Timestamp)
166+
check(assert_type(p3.to_timestamp("D", "E"), pd.Timestamp), pd.Timestamp)
167+
check(assert_type(p3.to_timestamp("D", "Finish"), pd.Timestamp), pd.Timestamp)
168+
check(assert_type(p3.to_timestamp("D", "End"), pd.Timestamp), pd.Timestamp)
169+
check(assert_type(p3.to_timestamp("D", "Begin"), pd.Timestamp), pd.Timestamp)
170+
check(assert_type(p3.to_timestamp("D", "Start"), pd.Timestamp), pd.Timestamp)
171+
172+
check(assert_type(p3.asfreq("D", "S"), pd.Period), pd.Period)
173+
check(assert_type(p3.asfreq(Day(), "E"), pd.Period), pd.Period)
174+
check(assert_type(p3.asfreq(Day(), "Finish"), pd.Period), pd.Period)
175+
check(assert_type(p3.asfreq(Day(), "Begin"), pd.Period), pd.Period)
176+
check(assert_type(p3.asfreq(Day(), "Start"), pd.Period), pd.Period)
177+
check(assert_type(p3.asfreq(Day(), "End"), pd.Period), pd.Period)
178+
check(assert_type(p3.asfreq(Day(), "end"), pd.Period), pd.Period)
179+
check(assert_type(p3.asfreq(Day(), "start"), pd.Period), pd.Period)
180+
check(assert_type(p3.asfreq(Day(), "begin"), pd.Period), pd.Period)
181+
check(assert_type(p3.asfreq(Day(), "finish"), pd.Period), pd.Period)
182+
check(assert_type(p3.asfreq(Day(), "s"), pd.Period), pd.Period)
183+
check(assert_type(p3.asfreq(Day(), "e"), pd.Period), pd.Period)
184+
185+
check(assert_type(pd.Period.now("D"), pd.Period), pd.Period)
186+
check(assert_type(pd.Period.now(Day()), pd.Period), pd.Period)
187+
188+
check(assert_type(p.strftime("%Y-%m-%d"), str), str)
189+
check(assert_type(hash(p), int), int)

0 commit comments

Comments
 (0)