Skip to content

Commit

Permalink
TYP: offsets.pyi (#45331)
Browse files Browse the repository at this point in the history
  • Loading branch information
twoertwein authored Feb 27, 2022
1 parent 7c59260 commit d391721
Show file tree
Hide file tree
Showing 28 changed files with 413 additions and 70 deletions.
2 changes: 2 additions & 0 deletions pandas/_libs/tslibs/ccalendar.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -8,5 +8,7 @@ def get_firstbday(year: int, month: int) -> int: ...
def get_lastbday(year: int, month: int) -> int: ...
def get_day_of_year(year: int, month: int, day: int) -> int: ...
def get_iso_calendar(year: int, month: int, day: int) -> tuple[int, int, int]: ...
def is_leapyear(year: int) -> bool: ...
def get_week_of_year(year: int, month: int, day: int) -> int: ...
def get_days_in_month(year: int, month: int) -> int: ...
def dayofweek(y: int, m: int, d: int) -> int: ...
256 changes: 256 additions & 0 deletions pandas/_libs/tslibs/offsets.pyi
Original file line number Diff line number Diff line change
@@ -0,0 +1,256 @@
from __future__ import annotations

from datetime import (
datetime,
timedelta,
)
from typing import (
TYPE_CHECKING,
Any,
Collection,
Literal,
Tuple,
TypeVar,
Union,
overload,
)

import numpy as np

from pandas._typing import npt

from .timedeltas import Timedelta

if TYPE_CHECKING:
from pandas.core.indexes.datetimes import DatetimeIndex
_BaseOffsetT = TypeVar("_BaseOffsetT", bound="BaseOffset")
_DatetimeT = TypeVar("_DatetimeT", bound=datetime)
_TimedeltaT = TypeVar("_TimedeltaT", bound=timedelta)

_relativedelta_kwds: set[str]
prefix_mapping: dict[str, type]

class ApplyTypeError(TypeError): ...

class BaseOffset:
n: int
def __init__(self, n: int = ..., normalize: bool = ...) -> None: ...
def __eq__(self, other) -> bool: ...
def __ne__(self, other) -> bool: ...
def __hash__(self) -> int: ...
@property
def kwds(self) -> dict: ...
@property
def base(self) -> BaseOffset: ...
@overload
def __add__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
@overload
def __add__(self: _BaseOffsetT, other: BaseOffset) -> _BaseOffsetT: ...
@overload
def __add__(self, other: _DatetimeT) -> _DatetimeT: ...
@overload
def __add__(self, other: _TimedeltaT) -> _TimedeltaT: ...
@overload
def __radd__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
@overload
def __radd__(self: _BaseOffsetT, other: BaseOffset) -> _BaseOffsetT: ...
@overload
def __radd__(self, other: _DatetimeT) -> _DatetimeT: ...
@overload
def __radd__(self, other: _TimedeltaT) -> _TimedeltaT: ...
def __sub__(self: _BaseOffsetT, other: BaseOffset) -> _BaseOffsetT: ...
@overload
def __rsub__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
@overload
def __rsub__(self: _BaseOffsetT, other: BaseOffset) -> _BaseOffsetT: ...
@overload
def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ...
@overload
def __rsub__(self, other: _TimedeltaT) -> _TimedeltaT: ...
def __call__(self, other): ...
@overload
def __mul__(self, other: np.ndarray) -> np.ndarray: ...
@overload
def __mul__(self: _BaseOffsetT, other: int) -> _BaseOffsetT: ...
@overload
def __rmul__(self, other: np.ndarray) -> np.ndarray: ...
@overload
def __rmul__(self: _BaseOffsetT, other: int) -> _BaseOffsetT: ...
def __neg__(self: _BaseOffsetT) -> _BaseOffsetT: ...
def copy(self: _BaseOffsetT) -> _BaseOffsetT: ...
def __repr__(self) -> str: ...
@property
def name(self) -> str: ...
@property
def rule_code(self) -> str: ...
def freqstr(self) -> str: ...
def apply_index(self, dtindex: "DatetimeIndex") -> "DatetimeIndex": ...
def _apply_array(self, dtarr) -> None: ...
def rollback(self, dt: datetime) -> datetime: ...
def rollforward(self, dt: datetime) -> datetime: ...
def is_on_offset(self, dt: datetime) -> bool: ...
def __setstate__(self, state) -> None: ...
def __getstate__(self): ...
@property
def nanos(self) -> int: ...
def onOffset(self, dt: datetime) -> bool: ...
def isAnchored(self) -> bool: ...
def is_anchored(self) -> bool: ...

def _get_offset(name: str) -> BaseOffset: ...

class SingleConstructorOffset(BaseOffset):
@classmethod
def _from_name(cls, suffix=...): ...
def __reduce__(self): ...

@overload
def to_offset(freq: None) -> None: ...
@overload
def to_offset(freq: timedelta | BaseOffset | str) -> BaseOffset: ...

class Tick(SingleConstructorOffset):
def __init__(self, n: int = ..., normalize: bool = ...) -> None: ...
@property
def delta(self) -> Timedelta: ...
@property
def nanos(self) -> int: ...

def delta_to_tick(delta: timedelta) -> Tick: ...

class Day(Tick): ...
class Hour(Tick): ...
class Minute(Tick): ...
class Second(Tick): ...
class Milli(Tick): ...
class Micro(Tick): ...
class Nano(Tick): ...

class RelativeDeltaOffset(BaseOffset):
def __init__(self, n: int = ..., normalize: bool = ..., **kwds: Any) -> None: ...

class BusinessMixin(SingleConstructorOffset):
def __init__(
self, n: int = ..., normalize: bool = ..., offset: timedelta = ...
): ...

class BusinessDay(BusinessMixin): ...

class BusinessHour(BusinessMixin):
def __init__(
self,
n: int = ...,
normalize: bool = ...,
start: str | Collection[str] = ...,
end: str | Collection[str] = ...,
offset: timedelta = ...,
): ...

class WeekOfMonthMixin(SingleConstructorOffset): ...

class YearOffset(SingleConstructorOffset):
def __init__(
self, n: int = ..., normalize: bool = ..., month: int | None = ...
): ...

class BYearEnd(YearOffset): ...
class BYearBegin(YearOffset): ...
class YearEnd(YearOffset): ...
class YearBegin(YearOffset): ...

class QuarterOffset(SingleConstructorOffset):
def __init__(
self, n: int = ..., normalize: bool = ..., startingMonth: int | None = ...
) -> None: ...

class BQuarterEnd(QuarterOffset): ...
class BQuarterBegin(QuarterOffset): ...
class QuarterEnd(QuarterOffset): ...
class QuarterBegin(QuarterOffset): ...
class MonthOffset(SingleConstructorOffset): ...
class MonthEnd(MonthOffset): ...
class MonthBegin(MonthOffset): ...
class BusinessMonthEnd(MonthOffset): ...
class BusinessMonthBegin(MonthOffset): ...

class SemiMonthOffset(SingleConstructorOffset):
def __init__(
self, n: int = ..., normalize: bool = ..., day_of_month: int | None = ...
) -> None: ...

class SemiMonthEnd(SemiMonthOffset): ...
class SemiMonthBegin(SemiMonthOffset): ...

class Week(SingleConstructorOffset):
def __init__(
self, n: int = ..., normalize: bool = ..., weekday: int | None = ...
) -> None: ...

class WeekOfMonth(WeekOfMonthMixin): ...
class LastWeekOfMonth(WeekOfMonthMixin): ...

class FY5253Mixin(SingleConstructorOffset):
def __init__(
self,
n: int = ...,
normalize: bool = ...,
weekday: int = ...,
startingMonth: int = ...,
variation: str = ...,
) -> None: ...

class FY5253(FY5253Mixin): ...
class FY5253Quarter(FY5253Mixin): ...
class Easter(SingleConstructorOffset): ...

class _CustomBusinessMonth(BusinessMixin):
def __init__(
self,
n: int = ...,
normalize: bool = ...,
offset: timedelta = ...,
holidays: None | list = ...,
): ...

class CustomBusinessDay(BusinessDay):
def __init__(
self,
n: int = ...,
normalize: bool = ...,
offset: timedelta = ...,
weekmask: str = ...,
): ...

class CustomBusinessHour(BusinessHour):
def __init__(
self,
n: int = ...,
normalize: bool = ...,
start: str = ...,
end: str = ...,
offset: timedelta = ...,
holidays: None | list = ...,
): ...

class CustomBusinessMonthEnd(_CustomBusinessMonth): ...
class CustomBusinessMonthBegin(_CustomBusinessMonth): ...
class DateOffset(RelativeDeltaOffset): ...

BDay = BusinessDay
BMonthEnd = BusinessMonthEnd
BMonthBegin = BusinessMonthBegin
CBMonthEnd = CustomBusinessMonthEnd
CBMonthBegin = CustomBusinessMonthBegin
CDay = CustomBusinessDay

def roll_qtrday(
other: datetime, n: int, month: int, day_opt: str, modby: int
) -> int: ...

INVALID_FREQ_ERR_MSG: Literal["Invalid frequency: {0}"]

def shift_months(
dtindex: npt.NDArray[np.int64], months: int, day_opt: str | None = ...
) -> npt.NDArray[np.int64]: ...

_offset_map: dict[str, BaseOffset]
7 changes: 4 additions & 3 deletions pandas/_libs/tslibs/period.pyi
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from datetime import timedelta
from typing import Literal

import numpy as np
Expand Down Expand Up @@ -33,7 +34,7 @@ def get_period_field_arr(
) -> npt.NDArray[np.int64]: ...
def from_ordinals(
values: npt.NDArray[np.int64], # const int64_t[:]
freq: Frequency,
freq: timedelta | BaseOffset | str,
) -> npt.NDArray[np.int64]: ...
def extract_ordinals(
values: npt.NDArray[np.object_],
Expand All @@ -59,7 +60,7 @@ class Period:
def __new__( # type: ignore[misc]
cls,
value=...,
freq: int | str | None = ...,
freq: int | str | BaseOffset | None = ...,
ordinal: int | None = ...,
year: int | None = ...,
month: int | None = ...,
Expand All @@ -82,7 +83,7 @@ class Period:
how: str = ...,
tz: Timezone | None = ...,
) -> Timestamp: ...
def asfreq(self, freq: str, how: str = ...) -> Period: ...
def asfreq(self, freq: str | BaseOffset, how: str = ...) -> Period: ...
@property
def freqstr(self) -> str: ...
@property
Expand Down
2 changes: 1 addition & 1 deletion pandas/_libs/tslibs/period.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -1795,7 +1795,7 @@ cdef class _Period(PeriodMixin):

Parameters
----------
freq : str
freq : str, BaseOffset
The desired frequency.
how : {'E', 'S', 'end', 'start'}, default 'end'
Start or end of the timespan.
Expand Down
19 changes: 17 additions & 2 deletions pandas/_libs/tslibs/timedeltas.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -67,10 +67,25 @@ class Timedelta(timedelta):
def __abs__(self) -> timedelta: ...
def __mul__(self, other: float) -> timedelta: ...
def __rmul__(self, other: float) -> timedelta: ...
@overload
# error: Signature of "__floordiv__" incompatible with supertype "timedelta"
@overload # type: ignore[override]
def __floordiv__(self, other: timedelta) -> int: ...
@overload
def __floordiv__(self, other: int) -> timedelta: ...
def __floordiv__(self, other: int | float) -> timedelta: ...
@overload
def __floordiv__(
self, other: npt.NDArray[np.timedelta64]
) -> npt.NDArray[np.intp]: ...
@overload
def __floordiv__(
self, other: npt.NDArray[np.number]
) -> npt.NDArray[np.timedelta64] | Timedelta: ...
@overload
def __rfloordiv__(self, other: timedelta | str) -> int: ...
@overload
def __rfloordiv__(self, other: None | NaTType) -> NaTType: ...
@overload
def __rfloordiv__(self, other: np.ndarray) -> npt.NDArray[np.timedelta64]: ...
@overload
def __truediv__(self, other: timedelta) -> float: ...
@overload
Expand Down
9 changes: 5 additions & 4 deletions pandas/_libs/tslibs/timestamps.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ from pandas._libs.tslibs import (
BaseOffset,
NaTType,
Period,
Tick,
Timedelta,
)

Expand Down Expand Up @@ -139,14 +140,14 @@ class Timestamp(datetime):
@overload # type: ignore[override]
def __add__(self, other: np.ndarray) -> np.ndarray: ...
@overload
# TODO: other can also be Tick (but it cannot be resolved)
def __add__(self: _DatetimeT, other: timedelta | np.timedelta64) -> _DatetimeT: ...
def __add__(
self: _DatetimeT, other: timedelta | np.timedelta64 | Tick
) -> _DatetimeT: ...
def __radd__(self: _DatetimeT, other: timedelta) -> _DatetimeT: ...
@overload # type: ignore
def __sub__(self, other: datetime) -> timedelta: ...
@overload
# TODO: other can also be Tick (but it cannot be resolved)
def __sub__(self, other: timedelta | np.timedelta64) -> datetime: ...
def __sub__(self, other: timedelta | np.timedelta64 | Tick) -> datetime: ...
def __hash__(self) -> int: ...
def weekday(self) -> int: ...
def isoweekday(self) -> int: ...
Expand Down
5 changes: 3 additions & 2 deletions pandas/core/arrays/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@

from pandas.core.arrays import (
DatetimeArray,
PeriodArray,
TimedeltaArray,
)

Expand Down Expand Up @@ -1325,7 +1326,7 @@ def __add__(self, other):
elif is_integer_dtype(other_dtype):
if not is_period_dtype(self.dtype):
raise integer_op_not_supported(self)
result = self._addsub_int_array(other, operator.add)
result = cast("PeriodArray", self)._addsub_int_array(other, operator.add)
else:
# Includes Categorical, other ExtensionArrays
# For PeriodDtype, if self is a TimedeltaArray and other is a
Expand Down Expand Up @@ -1385,7 +1386,7 @@ def __sub__(self, other):
elif is_integer_dtype(other_dtype):
if not is_period_dtype(self.dtype):
raise integer_op_not_supported(self)
result = self._addsub_int_array(other, operator.sub)
result = cast("PeriodArray", self)._addsub_int_array(other, operator.sub)
else:
# Includes ExtensionArrays, float_dtype
return NotImplemented
Expand Down
Loading

0 comments on commit d391721

Please sign in to comment.