Skip to content

TYP: offsets.pyi #45331

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Feb 27, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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: ...
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These operators are not "perfect" but they might be good enough for now. Based on the current annotations, __add__(self: Year, other: Second) is declared to return a Year but it actually returns Second

@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