Skip to content

Commit e62cc10

Browse files
bashtagebashtage
committed
TST: Add tests for Period
1 parent 290f951 commit e62cc10

File tree

3 files changed

+252
-38
lines changed

3 files changed

+252
-38
lines changed

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

Lines changed: 84 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,50 @@
1-
from datetime import tzinfo
2-
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+
DatetimeIndex,
11+
Index,
12+
PeriodIndex,
13+
Timedelta,
14+
)
15+
from typing_extensions import TypeAlias
16+
17+
from pandas._typing import npt
318

419
from .timestamps import Timestamp
520

621
class IncompatibleFrequency(ValueError): ...
722

823
from pandas._libs.tslibs.offsets import BaseOffset
924

25+
_PeriodAddSub: TypeAlias = Union[
26+
Timedelta, datetime.timedelta, np.timedelta64, np.int64, int
27+
]
28+
29+
_PeriodEqualityComparison: TypeAlias = Union[
30+
Period, datetime.datetime, datetime.date, Timestamp, np.datetime64, int, np.int64
31+
]
32+
33+
_PeriodFreqHow: TypeAlias = Literal[
34+
"S",
35+
"E",
36+
"Start",
37+
"Finish",
38+
"Begin",
39+
"End",
40+
"s",
41+
"e",
42+
"start",
43+
"finish",
44+
"begin",
45+
"end",
46+
]
47+
1048
class PeriodMixin:
1149
@property
1250
def end_time(self) -> Timestamp: ...
@@ -16,21 +54,8 @@ class PeriodMixin:
1654
class Period(PeriodMixin):
1755
def __init__(
1856
self,
19-
value: Any = ...,
20-
freqstr: Any = ...,
21-
ordinal: Any = ...,
22-
year: Any = ...,
23-
month: int = ...,
24-
quarter: Any = ...,
25-
day: int = ...,
26-
hour: int = ...,
27-
minute: int = ...,
28-
second: int = ...,
29-
) -> None: ...
30-
def __new__(
31-
cls,
32-
value: Period | str = ...,
33-
freq: int | str | BaseOffset | None = ...,
57+
value: Period | str | None = ...,
58+
freq: str | BaseOffset | None = ...,
3459
ordinal: int | None = ...,
3560
year: int | None = ...,
3661
month: int | None = ...,
@@ -39,20 +64,45 @@ class Period(PeriodMixin):
3964
hour: int | None = ...,
4065
minute: int | None = ...,
4166
second: int | None = ...,
42-
) -> Period: ...
43-
def __sub__(self, other) -> Period | BaseOffset: ...
44-
def __add__(self, other) -> Period: ...
45-
def __eq__(self, other) -> bool: ...
46-
def __ge__(self, other) -> bool: ...
47-
def __gt__(self, other) -> bool: ...
67+
) -> None: ...
68+
@overload
69+
def __sub__(self, other: _PeriodAddSub) -> Period: ...
70+
@overload
71+
def __sub__(self, other: Period) -> BaseOffset: ...
72+
@overload
73+
def __sub__(self, other: PeriodIndex) -> Index: ...
74+
@overload
75+
def __add__(self, other: _PeriodAddSub) -> Period: ...
76+
@overload
77+
def __add__(self, other: Index) -> Period: ...
78+
@overload # type: ignore[override]
79+
def __eq__(self, other: _PeriodEqualityComparison) -> bool: ...
80+
@overload
81+
def __eq__(self, other: PeriodIndex | DatetimeIndex) -> npt.NDArray[np.bool_]: ...
82+
@overload
83+
def __ge__(self, other: Period) -> bool: ...
84+
@overload
85+
def __ge__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
86+
@overload
87+
def __gt__(self, other: Period) -> bool: ...
88+
@overload
89+
def __gt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
4890
def __hash__(self) -> int: ...
49-
def __le__(self, other) -> bool: ...
50-
def __lt__(self, other) -> bool: ...
51-
def __ne__(self, other) -> bool: ...
52-
def __radd__(self, other) -> Period: ...
53-
def __reduce__(self, *args, **kwargs) -> Any: ... # what should this be?
54-
def __rsub__(self, other) -> Period: ...
55-
def __setstate__(self, *args, **kwargs) -> Any: ... # what should this be?
91+
@overload
92+
def __le__(self, other: Period) -> bool: ...
93+
@overload
94+
def __le__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
95+
@overload
96+
def __lt__(self, other: Period) -> bool: ...
97+
@overload
98+
def __lt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ...
99+
@overload # type: ignore[override]
100+
def __ne__(self, other: _PeriodEqualityComparison) -> bool: ...
101+
@overload
102+
def __ne__(self, other: PeriodIndex | DatetimeIndex) -> npt.NDArray[np.bool_]: ...
103+
# Ignored due to indecipherable error from mypy:
104+
# Forward operator "__add__" is not callable [misc]
105+
def __radd__(self, other: _PeriodAddSub) -> Period: ... # type: ignore[misc]
56106
@property
57107
def day(self) -> int: ...
58108
@property
@@ -66,7 +116,7 @@ class Period(PeriodMixin):
66116
@property
67117
def end_time(self) -> Timestamp: ...
68118
@property
69-
def freq(self) -> Any: ...
119+
def freq(self) -> BaseOffset: ...
70120
@property
71121
def freqstr(self) -> str: ...
72122
@property
@@ -99,14 +149,12 @@ class Period(PeriodMixin):
99149
def day_of_year(self) -> int: ...
100150
@property
101151
def day_of_week(self) -> int: ...
102-
def asfreq(self, freq: str | BaseOffset, how: str = ...) -> Period: ...
152+
def asfreq(self, freq: str | BaseOffset, how: _PeriodFreqHow = ...) -> Period: ...
103153
@classmethod
104-
def now(cls, freq: BaseOffset = ...) -> Period: ...
154+
def now(cls, freq: str | BaseOffset = ...) -> Period: ...
105155
def strftime(self, fmt: str) -> str: ...
106156
def to_timestamp(
107157
self,
108158
freq: str | BaseOffset | None = ...,
109-
how: str = ...,
110-
tz: str | tzinfo | None = ...,
159+
how: _PeriodFreqHow = ...,
111160
) -> Timestamp: ...
112-
def astype(self, dtype, copy: bool = ...): ...

‎tests/test_dtypes.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
import numpy as np
77
import pandas as pd
8+
import pyarrow as pa
89
from typing_extensions import assert_type
910

1011
from tests import check
@@ -114,6 +115,4 @@ def test_boolean_dtype() -> None:
114115

115116

116117
def test_arrow_dtype() -> None:
117-
import pyarrow as pa
118-
119118
check(assert_type(pd.ArrowDtype(pa.int64()), pd.ArrowDtype), pd.ArrowDtype)

‎tests/test_timefuncs.py

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,12 @@
1919
from pandas._libs import NaTType
2020
from pandas._libs.tslibs import BaseOffset
2121

22+
from pandas.tseries.offsets import (
23+
BusinessDay,
24+
CustomBusinessDay,
25+
Day,
26+
)
27+
2228
if TYPE_CHECKING:
2329
from pandas._typing import FulldatetimeDict
2430
else:
@@ -988,3 +994,164 @@ def test_timedelta_range() -> None:
988994
),
989995
pd.TimedeltaIndex,
990996
)
997+
998+
999+
def test_period() -> None:
1000+
p = pd.Period("2012-1-1", freq="D")
1001+
check(assert_type(p, pd.Period), pd.Period)
1002+
check(assert_type(pd.Period(p), pd.Period), pd.Period)
1003+
check(assert_type(pd.Period("2012-1-1", freq=Day()), pd.Period), pd.Period)
1004+
check(
1005+
assert_type(pd.Period(freq="D", year=2012, day=1, month=1), pd.Period),
1006+
pd.Period,
1007+
)
1008+
check(
1009+
assert_type(pd.Period(None, "D", year=2012, day=1, month=1), pd.Period),
1010+
pd.Period,
1011+
)
1012+
check(
1013+
assert_type(pd.Period(None, "D", 1, year=2012, day=1, month=1), pd.Period),
1014+
pd.Period,
1015+
)
1016+
check(
1017+
assert_type(
1018+
pd.Period(
1019+
freq="s", year=2012, month=1, day=1, hour=12, minute=30, second=45
1020+
),
1021+
pd.Period,
1022+
),
1023+
pd.Period,
1024+
)
1025+
check(assert_type(pd.Period(freq="Q", year=2012, quarter=2), pd.Period), pd.Period)
1026+
check(assert_type(p.day, int), int)
1027+
check(assert_type(p.day_of_week, int), int)
1028+
check(assert_type(p.day_of_year, int), int)
1029+
check(assert_type(p.dayofweek, int), int)
1030+
check(assert_type(p.dayofyear, int), int)
1031+
check(assert_type(p.days_in_month, int), int)
1032+
check(assert_type(p.daysinmonth, int), int)
1033+
check(assert_type(p.end_time, pd.Timestamp), pd.Timestamp)
1034+
check(assert_type(p.freqstr, str), str)
1035+
check(assert_type(p.hour, int), int)
1036+
check(assert_type(p.is_leap_year, bool), bool)
1037+
check(assert_type(p.minute, int), int)
1038+
check(assert_type(p.month, int), int)
1039+
check(assert_type(p.quarter, int), int)
1040+
check(assert_type(p.qyear, int), int)
1041+
check(assert_type(p.second, int), int)
1042+
check(assert_type(p.start_time, pd.Timestamp), pd.Timestamp)
1043+
check(assert_type(p.week, int), int)
1044+
check(assert_type(p.weekday, int), int)
1045+
check(assert_type(p.weekofyear, int), int)
1046+
check(assert_type(p.year, int), int)
1047+
check(assert_type(p.freq, BaseOffset), Day)
1048+
check(assert_type(p.ordinal, int), int)
1049+
1050+
p2 = pd.Period("2012-1-1", freq="2D")
1051+
check(assert_type(p2.freq, BaseOffset), Day)
1052+
1053+
as0 = pd.Timedelta(1, "D")
1054+
as1 = dt.timedelta(days=1)
1055+
as2 = np.timedelta64(1, "D")
1056+
as3 = np.int64(1)
1057+
as4 = int(1)
1058+
as5 = pd.period_range("2012-1-1", periods=10, freq="D")
1059+
as6 = pd.Period("2012-1-1", freq="D")
1060+
1061+
check(assert_type(p + as0, pd.Period), pd.Period)
1062+
check(assert_type(p + as1, pd.Period), pd.Period)
1063+
check(assert_type(p + as2, pd.Period), pd.Period)
1064+
check(assert_type(p + as3, pd.Period), pd.Period)
1065+
check(assert_type(p + as4, pd.Period), pd.Period)
1066+
1067+
check(assert_type(p - as0, pd.Period), pd.Period)
1068+
check(assert_type(p - as1, pd.Period), pd.Period)
1069+
check(assert_type(p - as2, pd.Period), pd.Period)
1070+
check(assert_type(p - as3, pd.Period), pd.Period)
1071+
check(assert_type(p - as4, pd.Period), pd.Period)
1072+
check(assert_type(p - as5, pd.Index), pd.Index)
1073+
check(assert_type(p - as6, BaseOffset), Day)
1074+
1075+
check(assert_type(as0 + p, pd.Period), pd.Period)
1076+
check(assert_type(as1 + p, pd.Period), pd.Period)
1077+
check(assert_type(as2 + p, pd.Period), pd.Period)
1078+
check(assert_type(as3 + p, pd.Period), pd.Period)
1079+
check(assert_type(as4 + p, pd.Period), pd.Period)
1080+
# TOOD: PeriodIndex should have a __sub__ with correct types, this op is valid
1081+
# and so the assert_type is skipped
1082+
check(as5 - p, pd.Index) # type: ignore[operator]
1083+
1084+
check(assert_type(p.__radd__(as0), pd.Period), pd.Period)
1085+
check(assert_type(p.__radd__(as1), pd.Period), pd.Period)
1086+
check(assert_type(p.__radd__(as2), pd.Period), pd.Period)
1087+
check(assert_type(p.__radd__(as3), pd.Period), pd.Period)
1088+
check(assert_type(p.__radd__(as4), pd.Period), pd.Period)
1089+
1090+
p - p
1091+
1092+
c0 = pd.Timestamp("2012-1-1")
1093+
c1 = dt.datetime(2012, 1, 1)
1094+
c2 = dt.date(2012, 1, 1)
1095+
c3 = np.datetime64(1, "ns")
1096+
c4 = 0
1097+
c5 = np.int64(0)
1098+
c6 = pd.Period("2012-1-1", freq="D")
1099+
c7 = pd.period_range("2012-1-1", periods=10, freq="D")
1100+
c8 = pd.date_range("2012-1-1", periods=10, freq="D")
1101+
1102+
check(assert_type(p == c0, bool), bool)
1103+
check(assert_type(p == c1, bool), bool)
1104+
check(assert_type(p == c2, bool), bool)
1105+
check(assert_type(p == c3, bool), bool)
1106+
check(assert_type(p == c4, bool), bool)
1107+
check(assert_type(p == c5, bool), bool)
1108+
check(assert_type(p == c7, npt.NDArray[np.bool_]), np.ndarray)
1109+
check(assert_type(p == c8, npt.NDArray[np.bool_]), np.ndarray)
1110+
1111+
check(assert_type(p != c0, bool), bool)
1112+
check(assert_type(p != c1, bool), bool)
1113+
check(assert_type(p != c2, bool), bool)
1114+
check(assert_type(p != c3, bool), bool)
1115+
check(assert_type(p != c4, bool), bool)
1116+
check(assert_type(p != c5, bool), bool)
1117+
check(assert_type(p != c7, npt.NDArray[np.bool_]), np.ndarray)
1118+
check(assert_type(p != c8, npt.NDArray[np.bool_]), np.ndarray)
1119+
1120+
check(assert_type(p > c6, bool), bool)
1121+
check(assert_type(p > c7, npt.NDArray[np.bool_]), np.ndarray)
1122+
1123+
check(assert_type(p < c6, bool), bool)
1124+
check(assert_type(p < c7, npt.NDArray[np.bool_]), np.ndarray)
1125+
1126+
check(assert_type(p <= c6, bool), bool)
1127+
check(assert_type(p <= c7, npt.NDArray[np.bool_]), np.ndarray)
1128+
1129+
check(assert_type(p >= c6, bool), bool)
1130+
check(assert_type(p >= c7, npt.NDArray[np.bool_]), np.ndarray)
1131+
1132+
p3 = pd.Period("2007-01", freq="M")
1133+
check(assert_type(p3.to_timestamp("D", "S"), pd.Timestamp), pd.Timestamp)
1134+
check(assert_type(p3.to_timestamp("D", "E"), pd.Timestamp), pd.Timestamp)
1135+
check(assert_type(p3.to_timestamp("D", "Finish"), pd.Timestamp), pd.Timestamp)
1136+
check(assert_type(p3.to_timestamp("D", "End"), pd.Timestamp), pd.Timestamp)
1137+
check(assert_type(p3.to_timestamp("D", "Begin"), pd.Timestamp), pd.Timestamp)
1138+
check(assert_type(p3.to_timestamp("D", "Start"), pd.Timestamp), pd.Timestamp)
1139+
1140+
check(assert_type(p3.asfreq("D", "S"), pd.Period), pd.Period)
1141+
check(assert_type(p3.asfreq(Day(), "E"), pd.Period), pd.Period)
1142+
check(assert_type(p3.asfreq(Day(), "Finish"), pd.Period), pd.Period)
1143+
check(assert_type(p3.asfreq(Day(), "Begin"), pd.Period), pd.Period)
1144+
check(assert_type(p3.asfreq(Day(), "Start"), pd.Period), pd.Period)
1145+
check(assert_type(p3.asfreq(Day(), "End"), pd.Period), pd.Period)
1146+
check(assert_type(p3.asfreq(Day(), "end"), pd.Period), pd.Period)
1147+
check(assert_type(p3.asfreq(Day(), "start"), pd.Period), pd.Period)
1148+
check(assert_type(p3.asfreq(Day(), "begin"), pd.Period), pd.Period)
1149+
check(assert_type(p3.asfreq(Day(), "finish"), pd.Period), pd.Period)
1150+
check(assert_type(p3.asfreq(Day(), "s"), pd.Period), pd.Period)
1151+
check(assert_type(p3.asfreq(Day(), "e"), pd.Period), pd.Period)
1152+
1153+
check(assert_type(pd.Period.now("D"), pd.Period), pd.Period)
1154+
check(assert_type(pd.Period.now(Day()), pd.Period), pd.Period)
1155+
1156+
check(assert_type(p.strftime("%Y-%m-%d"), str), str)
1157+
check(assert_type(hash(p), int), int)

0 commit comments

Comments
 (0)