Skip to content

Commit d391721

Browse files
twoertweintwoertwein
authored
TYP: offsets.pyi (#45331)
1 parent 7c59260 commit d391721

28 files changed

+413
-70
lines changed

‎pandas/_libs/tslibs/ccalendar.pyi

+2
Original file line numberDiff line numberDiff line change
@@ -8,5 +8,7 @@ def get_firstbday(year: int, month: int) -> int: ...
88
def get_lastbday(year: int, month: int) -> int: ...
99
def get_day_of_year(year: int, month: int, day: int) -> int: ...
1010
def get_iso_calendar(year: int, month: int, day: int) -> tuple[int, int, int]: ...
11+
def is_leapyear(year: int) -> bool: ...
1112
def get_week_of_year(year: int, month: int, day: int) -> int: ...
1213
def get_days_in_month(year: int, month: int) -> int: ...
14+
def dayofweek(y: int, m: int, d: int) -> int: ...

‎pandas/_libs/tslibs/offsets.pyi

+256
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,256 @@
1+
from __future__ import annotations
2+
3+
from datetime import (
4+
datetime,
5+
timedelta,
6+
)
7+
from typing import (
8+
TYPE_CHECKING,
9+
Any,
10+
Collection,
11+
Literal,
12+
Tuple,
13+
TypeVar,
14+
Union,
15+
overload,
16+
)
17+
18+
import numpy as np
19+
20+
from pandas._typing import npt
21+
22+
from .timedeltas import Timedelta
23+
24+
if TYPE_CHECKING:
25+
from pandas.core.indexes.datetimes import DatetimeIndex
26+
_BaseOffsetT = TypeVar("_BaseOffsetT", bound="BaseOffset")
27+
_DatetimeT = TypeVar("_DatetimeT", bound=datetime)
28+
_TimedeltaT = TypeVar("_TimedeltaT", bound=timedelta)
29+
30+
_relativedelta_kwds: set[str]
31+
prefix_mapping: dict[str, type]
32+
33+
class ApplyTypeError(TypeError): ...
34+
35+
class BaseOffset:
36+
n: int
37+
def __init__(self, n: int = ..., normalize: bool = ...) -> None: ...
38+
def __eq__(self, other) -> bool: ...
39+
def __ne__(self, other) -> bool: ...
40+
def __hash__(self) -> int: ...
41+
@property
42+
def kwds(self) -> dict: ...
43+
@property
44+
def base(self) -> BaseOffset: ...
45+
@overload
46+
def __add__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
47+
@overload
48+
def __add__(self: _BaseOffsetT, other: BaseOffset) -> _BaseOffsetT: ...
49+
@overload
50+
def __add__(self, other: _DatetimeT) -> _DatetimeT: ...
51+
@overload
52+
def __add__(self, other: _TimedeltaT) -> _TimedeltaT: ...
53+
@overload
54+
def __radd__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
55+
@overload
56+
def __radd__(self: _BaseOffsetT, other: BaseOffset) -> _BaseOffsetT: ...
57+
@overload
58+
def __radd__(self, other: _DatetimeT) -> _DatetimeT: ...
59+
@overload
60+
def __radd__(self, other: _TimedeltaT) -> _TimedeltaT: ...
61+
def __sub__(self: _BaseOffsetT, other: BaseOffset) -> _BaseOffsetT: ...
62+
@overload
63+
def __rsub__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
64+
@overload
65+
def __rsub__(self: _BaseOffsetT, other: BaseOffset) -> _BaseOffsetT: ...
66+
@overload
67+
def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ...
68+
@overload
69+
def __rsub__(self, other: _TimedeltaT) -> _TimedeltaT: ...
70+
def __call__(self, other): ...
71+
@overload
72+
def __mul__(self, other: np.ndarray) -> np.ndarray: ...
73+
@overload
74+
def __mul__(self: _BaseOffsetT, other: int) -> _BaseOffsetT: ...
75+
@overload
76+
def __rmul__(self, other: np.ndarray) -> np.ndarray: ...
77+
@overload
78+
def __rmul__(self: _BaseOffsetT, other: int) -> _BaseOffsetT: ...
79+
def __neg__(self: _BaseOffsetT) -> _BaseOffsetT: ...
80+
def copy(self: _BaseOffsetT) -> _BaseOffsetT: ...
81+
def __repr__(self) -> str: ...
82+
@property
83+
def name(self) -> str: ...
84+
@property
85+
def rule_code(self) -> str: ...
86+
def freqstr(self) -> str: ...
87+
def apply_index(self, dtindex: "DatetimeIndex") -> "DatetimeIndex": ...
88+
def _apply_array(self, dtarr) -> None: ...
89+
def rollback(self, dt: datetime) -> datetime: ...
90+
def rollforward(self, dt: datetime) -> datetime: ...
91+
def is_on_offset(self, dt: datetime) -> bool: ...
92+
def __setstate__(self, state) -> None: ...
93+
def __getstate__(self): ...
94+
@property
95+
def nanos(self) -> int: ...
96+
def onOffset(self, dt: datetime) -> bool: ...
97+
def isAnchored(self) -> bool: ...
98+
def is_anchored(self) -> bool: ...
99+
100+
def _get_offset(name: str) -> BaseOffset: ...
101+
102+
class SingleConstructorOffset(BaseOffset):
103+
@classmethod
104+
def _from_name(cls, suffix=...): ...
105+
def __reduce__(self): ...
106+
107+
@overload
108+
def to_offset(freq: None) -> None: ...
109+
@overload
110+
def to_offset(freq: timedelta | BaseOffset | str) -> BaseOffset: ...
111+
112+
class Tick(SingleConstructorOffset):
113+
def __init__(self, n: int = ..., normalize: bool = ...) -> None: ...
114+
@property
115+
def delta(self) -> Timedelta: ...
116+
@property
117+
def nanos(self) -> int: ...
118+
119+
def delta_to_tick(delta: timedelta) -> Tick: ...
120+
121+
class Day(Tick): ...
122+
class Hour(Tick): ...
123+
class Minute(Tick): ...
124+
class Second(Tick): ...
125+
class Milli(Tick): ...
126+
class Micro(Tick): ...
127+
class Nano(Tick): ...
128+
129+
class RelativeDeltaOffset(BaseOffset):
130+
def __init__(self, n: int = ..., normalize: bool = ..., **kwds: Any) -> None: ...
131+
132+
class BusinessMixin(SingleConstructorOffset):
133+
def __init__(
134+
self, n: int = ..., normalize: bool = ..., offset: timedelta = ...
135+
): ...
136+
137+
class BusinessDay(BusinessMixin): ...
138+
139+
class BusinessHour(BusinessMixin):
140+
def __init__(
141+
self,
142+
n: int = ...,
143+
normalize: bool = ...,
144+
start: str | Collection[str] = ...,
145+
end: str | Collection[str] = ...,
146+
offset: timedelta = ...,
147+
): ...
148+
149+
class WeekOfMonthMixin(SingleConstructorOffset): ...
150+
151+
class YearOffset(SingleConstructorOffset):
152+
def __init__(
153+
self, n: int = ..., normalize: bool = ..., month: int | None = ...
154+
): ...
155+
156+
class BYearEnd(YearOffset): ...
157+
class BYearBegin(YearOffset): ...
158+
class YearEnd(YearOffset): ...
159+
class YearBegin(YearOffset): ...
160+
161+
class QuarterOffset(SingleConstructorOffset):
162+
def __init__(
163+
self, n: int = ..., normalize: bool = ..., startingMonth: int | None = ...
164+
) -> None: ...
165+
166+
class BQuarterEnd(QuarterOffset): ...
167+
class BQuarterBegin(QuarterOffset): ...
168+
class QuarterEnd(QuarterOffset): ...
169+
class QuarterBegin(QuarterOffset): ...
170+
class MonthOffset(SingleConstructorOffset): ...
171+
class MonthEnd(MonthOffset): ...
172+
class MonthBegin(MonthOffset): ...
173+
class BusinessMonthEnd(MonthOffset): ...
174+
class BusinessMonthBegin(MonthOffset): ...
175+
176+
class SemiMonthOffset(SingleConstructorOffset):
177+
def __init__(
178+
self, n: int = ..., normalize: bool = ..., day_of_month: int | None = ...
179+
) -> None: ...
180+
181+
class SemiMonthEnd(SemiMonthOffset): ...
182+
class SemiMonthBegin(SemiMonthOffset): ...
183+
184+
class Week(SingleConstructorOffset):
185+
def __init__(
186+
self, n: int = ..., normalize: bool = ..., weekday: int | None = ...
187+
) -> None: ...
188+
189+
class WeekOfMonth(WeekOfMonthMixin): ...
190+
class LastWeekOfMonth(WeekOfMonthMixin): ...
191+
192+
class FY5253Mixin(SingleConstructorOffset):
193+
def __init__(
194+
self,
195+
n: int = ...,
196+
normalize: bool = ...,
197+
weekday: int = ...,
198+
startingMonth: int = ...,
199+
variation: str = ...,
200+
) -> None: ...
201+
202+
class FY5253(FY5253Mixin): ...
203+
class FY5253Quarter(FY5253Mixin): ...
204+
class Easter(SingleConstructorOffset): ...
205+
206+
class _CustomBusinessMonth(BusinessMixin):
207+
def __init__(
208+
self,
209+
n: int = ...,
210+
normalize: bool = ...,
211+
offset: timedelta = ...,
212+
holidays: None | list = ...,
213+
): ...
214+
215+
class CustomBusinessDay(BusinessDay):
216+
def __init__(
217+
self,
218+
n: int = ...,
219+
normalize: bool = ...,
220+
offset: timedelta = ...,
221+
weekmask: str = ...,
222+
): ...
223+
224+
class CustomBusinessHour(BusinessHour):
225+
def __init__(
226+
self,
227+
n: int = ...,
228+
normalize: bool = ...,
229+
start: str = ...,
230+
end: str = ...,
231+
offset: timedelta = ...,
232+
holidays: None | list = ...,
233+
): ...
234+
235+
class CustomBusinessMonthEnd(_CustomBusinessMonth): ...
236+
class CustomBusinessMonthBegin(_CustomBusinessMonth): ...
237+
class DateOffset(RelativeDeltaOffset): ...
238+
239+
BDay = BusinessDay
240+
BMonthEnd = BusinessMonthEnd
241+
BMonthBegin = BusinessMonthBegin
242+
CBMonthEnd = CustomBusinessMonthEnd
243+
CBMonthBegin = CustomBusinessMonthBegin
244+
CDay = CustomBusinessDay
245+
246+
def roll_qtrday(
247+
other: datetime, n: int, month: int, day_opt: str, modby: int
248+
) -> int: ...
249+
250+
INVALID_FREQ_ERR_MSG: Literal["Invalid frequency: {0}"]
251+
252+
def shift_months(
253+
dtindex: npt.NDArray[np.int64], months: int, day_opt: str | None = ...
254+
) -> npt.NDArray[np.int64]: ...
255+
256+
_offset_map: dict[str, BaseOffset]

‎pandas/_libs/tslibs/period.pyi

+4-3
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from datetime import timedelta
12
from typing import Literal
23

34
import numpy as np
@@ -33,7 +34,7 @@ def get_period_field_arr(
3334
) -> npt.NDArray[np.int64]: ...
3435
def from_ordinals(
3536
values: npt.NDArray[np.int64], # const int64_t[:]
36-
freq: Frequency,
37+
freq: timedelta | BaseOffset | str,
3738
) -> npt.NDArray[np.int64]: ...
3839
def extract_ordinals(
3940
values: npt.NDArray[np.object_],
@@ -59,7 +60,7 @@ class Period:
5960
def __new__( # type: ignore[misc]
6061
cls,
6162
value=...,
62-
freq: int | str | None = ...,
63+
freq: int | str | BaseOffset | None = ...,
6364
ordinal: int | None = ...,
6465
year: int | None = ...,
6566
month: int | None = ...,
@@ -82,7 +83,7 @@ class Period:
8283
how: str = ...,
8384
tz: Timezone | None = ...,
8485
) -> Timestamp: ...
85-
def asfreq(self, freq: str, how: str = ...) -> Period: ...
86+
def asfreq(self, freq: str | BaseOffset, how: str = ...) -> Period: ...
8687
@property
8788
def freqstr(self) -> str: ...
8889
@property

‎pandas/_libs/tslibs/period.pyx

+1-1
Original file line numberDiff line numberDiff line change
@@ -1795,7 +1795,7 @@ cdef class _Period(PeriodMixin):
17951795

17961796
Parameters
17971797
----------
1798-
freq : str
1798+
freq : str, BaseOffset
17991799
The desired frequency.
18001800
how : {'E', 'S', 'end', 'start'}, default 'end'
18011801
Start or end of the timespan.

‎pandas/_libs/tslibs/timedeltas.pyi

+17-2
Original file line numberDiff line numberDiff line change
@@ -67,10 +67,25 @@ class Timedelta(timedelta):
6767
def __abs__(self) -> timedelta: ...
6868
def __mul__(self, other: float) -> timedelta: ...
6969
def __rmul__(self, other: float) -> timedelta: ...
70-
@overload
70+
# error: Signature of "__floordiv__" incompatible with supertype "timedelta"
71+
@overload # type: ignore[override]
7172
def __floordiv__(self, other: timedelta) -> int: ...
7273
@overload
73-
def __floordiv__(self, other: int) -> timedelta: ...
74+
def __floordiv__(self, other: int | float) -> timedelta: ...
75+
@overload
76+
def __floordiv__(
77+
self, other: npt.NDArray[np.timedelta64]
78+
) -> npt.NDArray[np.intp]: ...
79+
@overload
80+
def __floordiv__(
81+
self, other: npt.NDArray[np.number]
82+
) -> npt.NDArray[np.timedelta64] | Timedelta: ...
83+
@overload
84+
def __rfloordiv__(self, other: timedelta | str) -> int: ...
85+
@overload
86+
def __rfloordiv__(self, other: None | NaTType) -> NaTType: ...
87+
@overload
88+
def __rfloordiv__(self, other: np.ndarray) -> npt.NDArray[np.timedelta64]: ...
7489
@overload
7590
def __truediv__(self, other: timedelta) -> float: ...
7691
@overload

‎pandas/_libs/tslibs/timestamps.pyi

+5-4
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ from pandas._libs.tslibs import (
1818
BaseOffset,
1919
NaTType,
2020
Period,
21+
Tick,
2122
Timedelta,
2223
)
2324

@@ -139,14 +140,14 @@ class Timestamp(datetime):
139140
@overload # type: ignore[override]
140141
def __add__(self, other: np.ndarray) -> np.ndarray: ...
141142
@overload
142-
# TODO: other can also be Tick (but it cannot be resolved)
143-
def __add__(self: _DatetimeT, other: timedelta | np.timedelta64) -> _DatetimeT: ...
143+
def __add__(
144+
self: _DatetimeT, other: timedelta | np.timedelta64 | Tick
145+
) -> _DatetimeT: ...
144146
def __radd__(self: _DatetimeT, other: timedelta) -> _DatetimeT: ...
145147
@overload # type: ignore
146148
def __sub__(self, other: datetime) -> timedelta: ...
147149
@overload
148-
# TODO: other can also be Tick (but it cannot be resolved)
149-
def __sub__(self, other: timedelta | np.timedelta64) -> datetime: ...
150+
def __sub__(self, other: timedelta | np.timedelta64 | Tick) -> datetime: ...
150151
def __hash__(self) -> int: ...
151152
def weekday(self) -> int: ...
152153
def isoweekday(self) -> int: ...

‎pandas/core/arrays/datetimelike.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@
130130

131131
from pandas.core.arrays import (
132132
DatetimeArray,
133+
PeriodArray,
133134
TimedeltaArray,
134135
)
135136

@@ -1325,7 +1326,7 @@ def __add__(self, other):
13251326
elif is_integer_dtype(other_dtype):
13261327
if not is_period_dtype(self.dtype):
13271328
raise integer_op_not_supported(self)
1328-
result = self._addsub_int_array(other, operator.add)
1329+
result = cast("PeriodArray", self)._addsub_int_array(other, operator.add)
13291330
else:
13301331
# Includes Categorical, other ExtensionArrays
13311332
# For PeriodDtype, if self is a TimedeltaArray and other is a
@@ -1385,7 +1386,7 @@ def __sub__(self, other):
13851386
elif is_integer_dtype(other_dtype):
13861387
if not is_period_dtype(self.dtype):
13871388
raise integer_op_not_supported(self)
1388-
result = self._addsub_int_array(other, operator.sub)
1389+
result = cast("PeriodArray", self)._addsub_int_array(other, operator.sub)
13891390
else:
13901391
# Includes ExtensionArrays, float_dtype
13911392
return NotImplemented

0 commit comments

Comments
 (0)