diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index df4e56304..c726cef62 100644 --- a/pandas-stubs/_libs/interval.pyi +++ b/pandas-stubs/_libs/interval.pyi @@ -123,28 +123,37 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __mul__(self: Interval[float], y: float) -> Interval[float]: ... @overload + def __mul__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ... + @overload def __rmul__( self: Interval[int], y: _OrderableScalarT ) -> Interval[_OrderableScalarT]: ... @overload def __rmul__(self: Interval[float], y: float) -> Interval[float]: ... @overload + def __rmul__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ... + @overload def __truediv__( self: Interval[int], y: _OrderableScalarT ) -> Interval[_OrderableScalarT]: ... @overload def __truediv__(self: Interval[float], y: float) -> Interval[float]: ... @overload + def __truediv__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ... + @overload def __floordiv__( self: Interval[int], y: _OrderableScalarT ) -> Interval[_OrderableScalarT]: ... @overload def __floordiv__(self: Interval[float], y: float) -> Interval[float]: ... + @overload + def __floordiv__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ... + @overload def overlaps(self: Interval[_OrderableT], other: Interval[_OrderableT]) -> bool: ... - -def intervals_to_interval_bounds( - intervals: np.ndarray, validate_closed: bool = ... -) -> tuple[np.ndarray, np.ndarray, str]: ... + @overload + def overlaps(self: Interval[int], other: Interval[float]) -> bool: ... + @overload + def overlaps(self: Interval[float], other: Interval[int]) -> bool: ... class IntervalTree(IntervalMixin): def __init__( diff --git a/pandas-stubs/_libs/tslibs/period.pyi b/pandas-stubs/_libs/tslibs/period.pyi index e0f93e54c..2b0d919f0 100644 --- a/pandas-stubs/_libs/tslibs/period.pyi +++ b/pandas-stubs/_libs/tslibs/period.pyi @@ -1,34 +1,120 @@ -from typing import Any +import datetime +from typing import ( + Literal, + Union, + overload, +) + +import numpy as np +from pandas import ( + Index, + PeriodIndex, + Timedelta, +) +from pandas.core.series import ( + PeriodSeries, + TimedeltaSeries, +) +from typing_extensions import TypeAlias + +from pandas._typing import npt + +from .timestamps import Timestamp class IncompatibleFrequency(ValueError): ... -class Period: +from pandas._libs.tslibs.offsets import BaseOffset + +_PeriodAddSub: TypeAlias = Union[ + Timedelta, datetime.timedelta, np.timedelta64, np.int64, int, BaseOffset +] + +_PeriodFreqHow: TypeAlias = Literal[ + "S", + "E", + "Start", + "Finish", + "Begin", + "End", + "s", + "e", + "start", + "finish", + "begin", + "end", +] + +class PeriodMixin: + @property + def end_time(self) -> Timestamp: ... + @property + def start_time(self) -> Timestamp: ... + +class Period(PeriodMixin): def __init__( self, - value: Any = ..., - freqstr: Any = ..., - ordinal: Any = ..., - year: Any = ..., - month: int = ..., - quarter: Any = ..., - day: int = ..., - hour: int = ..., - minute: int = ..., - second: int = ..., + value: Period | str | None = ..., + freq: str | BaseOffset | None = ..., + ordinal: int | None = ..., + year: int | None = ..., + month: int | None = ..., + quarter: int | None = ..., + day: int | None = ..., + hour: int | None = ..., + minute: int | None = ..., + second: int | None = ..., ) -> None: ... - def __add__(self, other) -> Period: ... - def __eq__(self, other) -> bool: ... - def __ge__(self, other) -> bool: ... - def __gt__(self, other) -> bool: ... + @overload + def __sub__(self, other: _PeriodAddSub) -> Period: ... + @overload + def __sub__(self, other: Period) -> BaseOffset: ... + @overload + def __sub__(self, other: PeriodIndex) -> Index: ... + @overload + def __sub__(self, other: TimedeltaSeries) -> PeriodSeries: ... + @overload + def __add__(self, other: _PeriodAddSub) -> Period: ... + @overload + def __add__(self, other: Index) -> PeriodIndex: ... + @overload + def __add__(self, other: TimedeltaSeries) -> PeriodSeries: ... + @overload # type: ignore[override] + def __eq__(self, other: Period) -> bool: ... + @overload + def __eq__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... + @overload + def __ge__(self, other: Period) -> bool: ... + @overload + def __ge__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... + @overload + def __gt__(self, other: Period) -> bool: ... + @overload + def __gt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... def __hash__(self) -> int: ... - def __le__(self, other) -> bool: ... - def __lt__(self, other) -> bool: ... - def __new__(cls, *args, **kwargs) -> Period: ... - def __ne__(self, other) -> bool: ... - def __radd__(self, other) -> Period: ... - def __reduce__(self, *args, **kwargs) -> Any: ... # what should this be? - def __rsub__(self, other) -> Period: ... - def __setstate__(self, *args, **kwargs) -> Any: ... # what should this be? + @overload + def __le__(self, other: Period) -> bool: ... + @overload + def __le__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... + @overload + def __lt__(self, other: Period) -> bool: ... + @overload + def __lt__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... + @overload # type: ignore[override] + def __ne__(self, other: Period) -> bool: ... + @overload + def __ne__(self, other: PeriodIndex) -> npt.NDArray[np.bool_]: ... + # Ignored due to indecipherable error from mypy: + # Forward operator "__add__" is not callable [misc] + @overload + def __radd__(self, other: _PeriodAddSub) -> Period: ... # type: ignore[misc] + # Real signature is -> PeriodIndex, but conflicts with Index.__add__ + # Changing Index is very hard due to Index inheritance + # Signatures of "__radd__" of "Period" and "__add__" of "Index" + # are unsafely overlapping + @overload + def __radd__(self, other: Index) -> Index: ... + @overload + def __radd__(self, other: TimedeltaSeries) -> PeriodSeries: ... @property def day(self) -> int: ... @property @@ -42,7 +128,7 @@ class Period: @property def end_time(self) -> Timestamp: ... @property - def freq(self) -> Any: ... + def freq(self) -> BaseOffset: ... @property def freqstr(self) -> str: ... @property @@ -71,12 +157,16 @@ class Period: def weekofyear(self) -> int: ... @property def year(self) -> int: ... - # Static methods + @property + def day_of_year(self) -> int: ... + @property + def day_of_week(self) -> int: ... + def asfreq(self, freq: str | BaseOffset, how: _PeriodFreqHow = ...) -> Period: ... @classmethod - def now(cls) -> Period: ... - # Methods - def asfreq(self, freq: str, how: str = ...) -> Period: ... + def now(cls, freq: str | BaseOffset = ...) -> Period: ... def strftime(self, fmt: str) -> str: ... - def to_timestamp(self, freq: str, how: str = ...) -> Timestamp: ... - -from .timestamps import Timestamp + def to_timestamp( + self, + freq: str | BaseOffset | None = ..., + how: _PeriodFreqHow = ..., + ) -> Timestamp: ... diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index aa09000ad..d197ee246 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -1,21 +1,43 @@ +import datetime as dt from datetime import timedelta from typing import ( ClassVar, Literal, + NamedTuple, TypeVar, Union, overload, ) import numpy as np +from pandas import ( + DatetimeIndex, + Float64Index, + Int64Index, + PeriodIndex, + Series, + TimedeltaIndex, +) +from pandas.core.series import TimedeltaSeries from typing_extensions import TypeAlias from pandas._libs.tslibs import ( + BaseOffset, NaTType, - Tick, ) +from pandas._libs.tslibs.period import Period +from pandas._libs.tslibs.timestamps import Timestamp from pandas._typing import npt +class Components(NamedTuple): + days: int + hours: int + minutes: int + seconds: int + milliseconds: int + microseconds: int + nanoseconds: int + # This should be kept consistent with the keys in the dict timedelta_abbrevs # in pandas/_libs/tslibs/timedeltas.pyx TimeDeltaUnitChoices: TypeAlias = Literal[ @@ -70,28 +92,23 @@ UnitChoices: TypeAlias = Union[ _S = TypeVar("_S", bound=timedelta) -def ints_to_pytimedelta( - arr: npt.NDArray[np.int64], # const int64_t[:] - box: bool = ..., -) -> npt.NDArray[np.object_]: ... -def array_to_timedelta64( - values: npt.NDArray[np.object_], - unit: str | None = ..., - errors: str = ..., -) -> np.ndarray: ... # np.ndarray[m8ns] -def parse_timedelta_unit(unit: str | None) -> UnitChoices: ... -def delta_to_nanoseconds(delta: np.timedelta64 | timedelta | Tick) -> int: ... - class Timedelta(timedelta): min: ClassVar[Timedelta] max: ClassVar[Timedelta] resolution: ClassVar[Timedelta] - value: int # np.int64 + value: int def __new__( cls: type[_S], - value=..., - unit: str = ..., - **kwargs: float | np.integer | np.floating, + value: str | int | Timedelta | timedelta | np.timedelta64 = ..., + unit: TimeDeltaUnitChoices = ..., + *, + days: float | np.integer | np.floating = ..., + seconds: float | np.integer | np.floating = ..., + microseconds: float | np.integer | np.floating = ..., + milliseconds: float | np.integer | np.floating = ..., + minutes: float | np.integer | np.floating = ..., + hours: float | np.integer | np.floating = ..., + weeks: float | np.integer | np.floating = ..., ) -> _S: ... # GH 46171 # While Timedelta can return pd.NaT, having the constructor return @@ -99,6 +116,8 @@ class Timedelta(timedelta): @property def days(self) -> int: ... @property + def nanoseconds(self) -> int: ... + @property def seconds(self) -> int: ... @property def microseconds(self) -> int: ... @@ -108,49 +127,177 @@ class Timedelta(timedelta): @property def asm8(self) -> np.timedelta64: ... # TODO: round/floor/ceil could return NaT? - def round(self: _S, freq: str) -> _S: ... - def floor(self: _S, freq: str) -> _S: ... - def ceil(self: _S, freq: str) -> _S: ... + def round(self: _S, freq: str | BaseOffset) -> _S: ... + def floor(self: _S, freq: str | BaseOffset) -> _S: ... + def ceil(self: _S, freq: str | BaseOffset) -> _S: ... @property def resolution_string(self) -> str: ... - def __add__(self, other: timedelta) -> Timedelta: ... + @overload # type: ignore[override] + def __add__(self, other: timedelta | np.timedelta64) -> Timedelta: ... + @overload + def __add__(self, other: dt.datetime | np.datetime64 | Timestamp) -> Timestamp: ... + @overload + def __add__(self, other: Period) -> Period: ... + @overload + def __add__(self, other: dt.date) -> dt.date: ... + @overload + def __add__(self, other: PeriodIndex) -> PeriodIndex: ... + @overload + def __add__(self, other: DatetimeIndex) -> DatetimeIndex: ... + @overload + def __add__( + self, other: npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.timedelta64]: ... + @overload + def __add__( + self, other: npt.NDArray[np.datetime64] + ) -> npt.NDArray[np.datetime64]: ... + @overload def __radd__(self, other: timedelta) -> Timedelta: ... - def __sub__(self, other: timedelta) -> Timedelta: ... - def __rsub__(self, other: timedelta) -> Timedelta: ... + @overload + def __radd__( + self, other: npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.timedelta64]: ... + @overload + def __radd__( + self, other: npt.NDArray[np.datetime64] + ) -> npt.NDArray[np.datetime64]: ... + @overload # type: ignore[override] + def __sub__(self, other: timedelta | np.timedelta64) -> Timedelta: ... + @overload + def __sub__( + self, other: npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.timedelta64]: ... + @overload + def __rsub__(self, other: timedelta | np.timedelta64) -> Timedelta: ... + @overload + def __rsub__(self, other: Timestamp | np.datetime64) -> Timestamp: ... + @overload + def __rsub__(self, other: PeriodIndex) -> PeriodIndex: ... + @overload + def __rsub__(self, other: DatetimeIndex) -> DatetimeIndex: ... + @overload + def __rsub__( + self, other: npt.NDArray[np.datetime64] + ) -> npt.NDArray[np.datetime64]: ... + @overload + def __rsub__( + self, other: npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.timedelta64]: ... def __neg__(self) -> Timedelta: ... def __pos__(self) -> Timedelta: ... def __abs__(self) -> Timedelta: ... + @overload # type: ignore[override] def __mul__(self, other: float) -> Timedelta: ... + @overload + def __mul__(self, other: np.ndarray) -> np.ndarray: ... + @overload + def __mul__(self, other: Series) -> Series: ... + @overload + def __mul__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... + @overload def __rmul__(self, other: float) -> Timedelta: ... + @overload + def __rmul__(self, other: np.ndarray) -> np.ndarray: ... + @overload + def __rmul__(self, other: Series) -> Series: ... + @overload + def __rmul__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... # error: Signature of "__floordiv__" incompatible with supertype "timedelta" @overload # type: ignore[override] def __floordiv__(self, other: timedelta) -> int: ... @overload def __floordiv__(self, other: float) -> Timedelta: ... @overload + def __floordiv__( + self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating] + ) -> npt.NDArray[np.timedelta64]: ... + @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: ... + def __floordiv__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... + @overload + def __floordiv__(self, other: Series) -> Series: ... @overload def __rfloordiv__(self, other: timedelta | str) -> int: ... @overload def __rfloordiv__(self, other: NaTType | None) -> NaTType: ... @overload def __rfloordiv__(self, other: np.ndarray) -> npt.NDArray[np.timedelta64]: ... - @overload + @overload # type: ignore[override] def __truediv__(self, other: timedelta) -> float: ... @overload def __truediv__(self, other: float) -> Timedelta: ... + @overload + def __truediv__( + self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating] + ) -> npt.NDArray[np.timedelta64]: ... + @overload + def __truediv__(self, other: TimedeltaSeries) -> Series[float]: ... + @overload + def __truediv__(self, other: Series) -> Series: ... + @overload + def __truediv__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... + @overload # type: ignore[override] + def __eq__(self, other: timedelta | np.timedelta64) -> bool: ... + @overload + def __eq__(self, other: Series) -> Series: ... + @overload + def __eq__( + self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.bool_]: ... + @overload # type: ignore[override] + def __ne__(self, other: timedelta | np.timedelta64) -> bool: ... + @overload + def __ne__(self, other: Series) -> Series: ... + @overload + def __ne__( + self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.bool_]: ... + @overload # type: ignore[override] def __mod__(self, other: timedelta) -> Timedelta: ... + @overload + def __mod__(self, other: float) -> Timedelta: ... + @overload + def __mod__( + self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating] + ) -> npt.NDArray[np.timedelta64]: ... + @overload + def __mod__(self, other: Series) -> Series: ... + @overload + def __mod__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... def __divmod__(self, other: timedelta) -> tuple[int, Timedelta]: ... - def __le__(self, other: timedelta) -> bool: ... - def __lt__(self, other: timedelta) -> bool: ... - def __ge__(self, other: timedelta) -> bool: ... - def __gt__(self, other: timedelta) -> bool: ... + # Mypy complains Forward operator "" is not callable, so ignore misc + # for le, lt ge and gt + @overload # type: ignore[override] + def __le__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] + @overload + def __le__( + self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.bool_]: ... + @overload # type: ignore[override] + def __lt__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] + @overload + def __lt__( + self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.bool_]: ... + @overload # type: ignore[override] + def __ge__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] + @overload + def __ge__( + self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.bool_]: ... + @overload # type: ignore[override] + def __gt__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] + @overload + def __gt__( + self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.bool_]: ... def __hash__(self) -> int: ... def isoformat(self) -> str: ... def to_numpy(self) -> np.timedelta64: ... + @property + def components(self) -> Components: ... + def view(self, dtype: npt.DTypeLike = ...) -> object: ... diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi index 17ed5f211..948c4ac0a 100644 --- a/pandas-stubs/_libs/tslibs/timestamps.pyi +++ b/pandas-stubs/_libs/tslibs/timestamps.pyi @@ -8,14 +8,20 @@ from datetime import ( from time import struct_time from typing import ( ClassVar, + Literal, TypeVar, overload, ) import numpy as np -from pandas import Index +from pandas import ( + DatetimeIndex, + Index, + TimedeltaIndex, +) from pandas.core.series import ( Series, + TimedeltaSeries, TimestampSeries, ) @@ -25,21 +31,24 @@ from pandas._libs.tslibs import ( Tick, Timedelta, ) -from pandas._typing import np_ndarray_bool +from pandas._typing import ( + np_ndarray_bool, + npt, +) _DatetimeT = TypeVar("_DatetimeT", bound=datetime) -def integer_op_not_supported(obj: object) -> TypeError: ... - class Timestamp(datetime): min: ClassVar[Timestamp] max: ClassVar[Timestamp] resolution: ClassVar[Timedelta] - value: int # np.int64 + value: int def __new__( cls: type[_DatetimeT], ts_input: np.integer | float | str | _date | datetime | np.datetime64 = ..., + # Freq is deprecated but is left in to allow code like Timestamp(2000,1,1) + # Removing it would make the other arguments position only freq: int | str | BaseOffset | None = ..., tz: str | _tzinfo | int | None = ..., unit: str | int | None = ..., @@ -52,8 +61,7 @@ class Timestamp(datetime): microsecond: int | None = ..., nanosecond: int | None = ..., tzinfo: _tzinfo | None = ..., - *, - fold: int | None = ..., + fold: Literal[0, 1] | None = ..., ) -> _DatetimeT: ... # GH 46171 # While Timestamp can return pd.NaT, having the constructor return @@ -73,6 +81,8 @@ class Timestamp(datetime): @property def microsecond(self) -> int: ... @property + def nanosecond(self) -> int: ... + @property def tzinfo(self) -> _tzinfo | None: ... @property def tz(self) -> _tzinfo | None: ... @@ -113,15 +123,15 @@ class Timestamp(datetime): def timetz(self) -> _time: ... def replace( self, - year: int = ..., - month: int = ..., - day: int = ..., - hour: int = ..., - minute: int = ..., - second: int = ..., - microsecond: int = ..., + year: int | None = ..., + month: int | None = ..., + day: int | None = ..., + hour: int | None = ..., + minute: int | None = ..., + second: int | None = ..., + microsecond: int | None = ..., tzinfo: _tzinfo | None = ..., - fold: int = ..., + fold: int | None = ..., ) -> Timestamp: ... def astimezone(self: _DatetimeT, tz: _tzinfo | None = ...) -> _DatetimeT: ... def ctime(self) -> str: ... @@ -131,44 +141,75 @@ class Timestamp(datetime): def utcoffset(self) -> timedelta | None: ... def tzname(self) -> str | None: ... def dst(self) -> timedelta | None: ... + # Mypy complains Forward operator "" is not callable, so ignore misc + # for le, lt ge and gt @overload # type: ignore[override] - def __le__(self, other: datetime) -> bool: ... + def __le__(self, other: datetime | np.datetime64) -> bool: ... # type: ignore[misc] @overload - def __le__(self, other: Index) -> np_ndarray_bool: ... + def __le__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ... @overload def __le__(self, other: TimestampSeries) -> Series[bool]: ... @overload # type: ignore[override] - def __lt__(self, other: datetime) -> bool: ... + def __lt__(self, other: datetime | np.datetime64) -> bool: ... # type: ignore[misc] @overload - def __lt__(self, other: Index) -> np_ndarray_bool: ... + def __lt__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ... @overload def __lt__(self, other: TimestampSeries) -> Series[bool]: ... @overload # type: ignore[override] - def __ge__(self, other: datetime) -> bool: ... + def __ge__(self, other: datetime | np.datetime64) -> bool: ... # type: ignore[misc] @overload - def __ge__(self, other: Index) -> np_ndarray_bool: ... + def __ge__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ... @overload def __ge__(self, other: TimestampSeries) -> Series[bool]: ... @overload # type: ignore[override] - def __gt__(self, other: datetime) -> bool: ... + def __gt__(self, other: datetime | np.datetime64) -> bool: ... # type: ignore[misc] @overload - def __gt__(self, other: Index) -> np_ndarray_bool: ... + def __gt__(self, other: Index | npt.NDArray[np.datetime64]) -> np_ndarray_bool: ... @overload def __gt__(self, other: TimestampSeries) -> Series[bool]: ... # error: Signature of "__add__" incompatible with supertype "date"/"datetime" @overload # type: ignore[override] - def __add__(self, other: np.ndarray) -> np.ndarray: ... + def __add__( + self, other: npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.datetime64]: ... @overload def __add__( self: _DatetimeT, other: timedelta | np.timedelta64 | Tick ) -> _DatetimeT: ... + @overload + def __add__(self, other: Series) -> TimestampSeries: ... + @overload + def __add__(self, other: TimedeltaIndex) -> DatetimeIndex: ... + @overload def __radd__(self: _DatetimeT, other: timedelta) -> _DatetimeT: ... + @overload + def __radd__(self, other: TimedeltaIndex) -> DatetimeIndex: ... @overload # type: ignore[override] def __sub__(self, other: datetime) -> Timedelta: ... @overload def __sub__( self: _DatetimeT, other: timedelta | np.timedelta64 | Tick ) -> _DatetimeT: ... + @overload + def __sub__(self, other: TimedeltaIndex) -> DatetimeIndex: ... + @overload + def __sub__(self, other: TimedeltaSeries) -> TimestampSeries: ... + @overload + def __sub__( + self, other: npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.datetime64]: ... + @overload # type: ignore[override] + def __eq__(self, other: Timestamp | np.datetime64 | datetime) -> bool: ... + @overload + def __eq__(self, other: TimestampSeries) -> Series[bool]: ... + @overload + def __eq__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ... + @overload # type: ignore[override] + def __ne__(self, other: Timestamp | np.datetime64 | datetime) -> bool: ... + @overload + def __ne__(self, other: TimestampSeries) -> Series[bool]: ... + @overload + def __ne__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ... def __hash__(self) -> int: ... def weekday(self) -> int: ... def isoweekday(self) -> int: ... @@ -223,6 +264,8 @@ class Timestamp(datetime): @property def dayofyear(self) -> int: ... @property + def weekofyear(self) -> int: ... + @property def quarter(self) -> int: ... @property def week(self) -> int: ... diff --git a/pandas-stubs/core/arrays/boolean.pyi b/pandas-stubs/core/arrays/boolean.pyi index 935948cc6..5a71d3843 100644 --- a/pandas-stubs/core/arrays/boolean.pyi +++ b/pandas-stubs/core/arrays/boolean.pyi @@ -1,4 +1,5 @@ import numpy as np +from pandas.core.arrays import ExtensionArray from pandas._typing import ( Scalar, @@ -10,27 +11,23 @@ from pandas.core.dtypes.base import ExtensionDtype as ExtensionDtype from .masked import BaseMaskedArray as BaseMaskedArray class BooleanDtype(ExtensionDtype): - name: str = ... @property def na_value(self) -> Scalar: ... @property def type(self) -> type_t: ... @property def kind(self) -> str: ... - @classmethod - def construct_array_type(cls) -> type_t[BooleanArray]: ... def __from_arrow__(self, array): ... -def coerce_to_array(values, mask=..., copy: bool = ...): ... - class BooleanArray(BaseMaskedArray): def __init__( self, values: np.ndarray, mask: np.ndarray, copy: bool = ... ) -> None: ... + def __setitem__(self, key: int | np.ndarray | slice, value: object) -> None: ... @property def dtype(self): ... - def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): ... - def __setitem__(self, key, value) -> None: ... - def astype(self, dtype, copy: bool = ...): ... - def any(self, skipna: bool = ..., **kwargs): ... - def all(self, skipna: bool = ..., **kwargs): ... + def astype( + self, dtype: str | np.dtype, copy: bool = ... + ) -> np.ndarray | ExtensionArray: ... + def any(self, skipna: bool = ..., **kwargs) -> bool: ... + def all(self, skipna: bool = ..., **kwargs) -> bool: ... diff --git a/pandas-stubs/core/arrays/interval.pyi b/pandas-stubs/core/arrays/interval.pyi index 1b63ed724..c2036b19d 100644 --- a/pandas-stubs/core/arrays/interval.pyi +++ b/pandas-stubs/core/arrays/interval.pyi @@ -17,13 +17,17 @@ class IntervalArray(IntervalMixin, ExtensionArray): cls, data, closed=..., dtype=..., copy: bool = ..., verify_integrity: bool = ... ): ... @classmethod - def from_breaks(cls, breaks, closed: str = ..., copy: bool = ..., dtype=...): ... + def from_breaks( + cls, breaks, closed: str = ..., copy: bool = ..., dtype=... + ) -> IntervalArray: ... @classmethod def from_arrays( cls, left, right, closed: str = ..., copy: bool = ..., dtype=... - ): ... + ) -> IntervalArray: ... @classmethod - def from_tuples(cls, data, closed: str = ..., copy: bool = ..., dtype=...): ... + def from_tuples( + cls, data, closed: str = ..., copy: bool = ..., dtype=... + ) -> IntervalArray: ... def __iter__(self): ... def __len__(self) -> int: ... def __getitem__(self, value): ... diff --git a/pandas-stubs/core/arrays/period.pyi b/pandas-stubs/core/arrays/period.pyi index 081baeb94..ea0fac180 100644 --- a/pandas-stubs/core/arrays/period.pyi +++ b/pandas-stubs/core/arrays/period.pyi @@ -6,10 +6,11 @@ from pandas.core.arrays.datetimelike import ( DatetimeLikeArrayMixin, ) -from pandas._libs.tslibs import Timestamp -from pandas._libs.tslibs.period import Period as Period +from pandas._libs.tslibs import Period +from pandas._libs.tslibs.timestamps import Timestamp +from pandas._typing import npt as npt -from pandas.tseries.offsets import Tick as Tick +from pandas.tseries.offsets import Tick class PeriodArray(DatetimeLikeArrayMixin, DatelikeOps): __array_priority__: int = ... @@ -43,11 +44,8 @@ class PeriodArray(DatetimeLikeArrayMixin, DatelikeOps): def asfreq(self, freq: str | None = ..., how: str = ...) -> Period: ... def astype(self, dtype, copy: bool = ...): ... -def raise_on_incompatible(left, right): ... def period_array( data: Sequence[Period | None], freq: str | Tick | None = ..., copy: bool = ..., ) -> PeriodArray: ... -def validate_dtype_freq(dtype, freq): ... -def dt64arr_to_periodarr(data, freq, tz=...): ... diff --git a/pandas-stubs/core/arrays/sparse/array.pyi b/pandas-stubs/core/arrays/sparse/array.pyi index c447b97a5..eb0650919 100644 --- a/pandas-stubs/core/arrays/sparse/array.pyi +++ b/pandas-stubs/core/arrays/sparse/array.pyi @@ -5,6 +5,8 @@ from pandas.core.arrays import ( ) from pandas.core.base import PandasObject +from pandas._libs.sparse import SparseIndex + class SparseArray(PandasObject, ExtensionArray, ExtensionOpsMixin): def __init__( self, @@ -20,9 +22,9 @@ class SparseArray(PandasObject, ExtensionArray, ExtensionOpsMixin): def __array__(self, dtype=..., copy=...) -> np.ndarray: ... def __setitem__(self, key, value) -> None: ... @property - def sp_index(self): ... + def sp_index(self) -> SparseIndex: ... @property - def sp_values(self): ... + def sp_values(self) -> np.ndarray: ... @property def dtype(self): ... @property diff --git a/pandas-stubs/core/arrays/sparse/dtype.pyi b/pandas-stubs/core/arrays/sparse/dtype.pyi index 630b5be03..19d6fb163 100644 --- a/pandas-stubs/core/arrays/sparse/dtype.pyi +++ b/pandas-stubs/core/arrays/sparse/dtype.pyi @@ -24,10 +24,3 @@ class SparseDtype(ExtensionDtype): def subtype(self): ... @property def name(self): ... - @classmethod - def construct_array_type(cls): ... - @classmethod - def construct_from_string(cls, string): ... - @classmethod - def is_dtype(cls, dtype): ... - def update_dtype(self, dtype): ... diff --git a/pandas-stubs/core/arrays/string_.pyi b/pandas-stubs/core/arrays/string_.pyi index 700510120..6bb335a89 100644 --- a/pandas-stubs/core/arrays/string_.pyi +++ b/pandas-stubs/core/arrays/string_.pyi @@ -1,20 +1,32 @@ +from typing import ( + Literal, + Sequence, +) + +import numpy as np +import pandas as pd from pandas.core.arrays import PandasArray -from pandas._typing import type_t +from pandas._typing import ( + AnyArrayLike, + type_t, +) from pandas.core.dtypes.base import ExtensionDtype class StringDtype(ExtensionDtype): - name: str = ... - na_value = ... + def __init__(self, storage: Literal["python", "pyarrow"] | None) -> None: ... @property def type(self) -> type_t: ... - @classmethod - def construct_array_type(cls) -> type_t[StringArray]: ... def __from_arrow__(self, array): ... class StringArray(PandasArray): - def __init__(self, values, copy: bool = ...) -> None: ... + def __init__( + self, + # Also pd.NA and np.nan but not possible it seems + values: AnyArrayLike | Sequence[str | None], + copy: bool = ..., + ) -> None: ... def __arrow_array__(self, type=...): ... def __setitem__(self, key, value) -> None: ... def fillna(self, value=..., method=..., limit=...): ... diff --git a/pandas-stubs/core/construction.pyi b/pandas-stubs/core/construction.pyi index 3cd517da3..02d855ac5 100644 --- a/pandas-stubs/core/construction.pyi +++ b/pandas-stubs/core/construction.pyi @@ -1,20 +1,30 @@ -from typing import Sequence +from typing import ( + Sequence, + Union, +) import numpy as np +import pandas as pd from pandas.core.indexes.api import Index from pandas.core.series import Series from pandas._typing import ( ArrayLike, Dtype, + npt, ) from pandas.core.dtypes.dtypes import ExtensionDtype from pandas.core.dtypes.generic import ABCExtensionArray def array( - data: Sequence[object], - dtype: str | np.dtype | ExtensionDtype | None = ..., + # str is forbidden even though Sequence[object] allows "abc" + data: npt.NDArray | Sequence[object] | pd.Index | pd.Series, + dtype: str + | np.dtype[np.generic] + | ExtensionDtype + | type[Union[str, bool, float, int]] + | None = ..., copy: bool = ..., ) -> ABCExtensionArray: ... def extract_array(obj, extract_numpy: bool = ...): ... diff --git a/pandas-stubs/core/dtypes/dtypes.pyi b/pandas-stubs/core/dtypes/dtypes.pyi index 15651cd91..317c40dec 100644 --- a/pandas-stubs/core/dtypes/dtypes.pyi +++ b/pandas-stubs/core/dtypes/dtypes.pyi @@ -1,15 +1,15 @@ -from typing import ( - Any, - Sequence, -) +import datetime as dt +from typing import Any +import numpy as np from pandas.core.indexes.base import Index +from pandas.core.series import Series -from pandas._libs.tslibs import ( # , timezones as timezones - Period as Period, - Timestamp, +from pandas._libs.tslibs import BaseOffset +from pandas._typing import ( + Ordered, + npt, ) -from pandas._typing import Ordered from .base import ExtensionDtype as ExtensionDtype @@ -32,94 +32,50 @@ class PandasExtensionDtype(ExtensionDtype): @classmethod def reset_cache(cls) -> None: ... -class CategoricalDtypeType(type): ... - class CategoricalDtype(PandasExtensionDtype, ExtensionDtype): - name: _str = ... - type: type[CategoricalDtypeType] = ... - kind: _str = ... - str: _str = ... - base = ... def __init__( - self, categories: Sequence[Any] | None = ..., ordered: Ordered = ... + self, + categories: Series | Index | list[Any] | None = ..., + ordered: Ordered = ..., ) -> None: ... - @classmethod - def construct_from_string(cls, string: _str) -> CategoricalDtype: ... def __hash__(self) -> int: ... def __eq__(self, other) -> bool: ... - @classmethod - def construct_array_type(cls): ... - @staticmethod - def validate_ordered(ordered: Ordered) -> None: ... - @staticmethod - def validate_categories(categories, fastpath: bool = ...): ... - def update_dtype(self, dtype: _str | CategoricalDtype) -> CategoricalDtype: ... @property def categories(self) -> Index: ... @property def ordered(self) -> Ordered: ... class DatetimeTZDtype(PandasExtensionDtype): - type: type[Timestamp] = ... - kind: _str = ... - str: _str = ... - num: int = ... - base = ... - na_value = ... - def __init__(self, unit: _str = ..., tz=...) -> None: ... + def __init__( + self, unit: _str = ..., tz: str | int | dt.tzinfo | None = ... + ) -> None: ... @property def unit(self): ... @property def tz(self): ... - @classmethod - def construct_array_type(cls): ... - @classmethod - def construct_from_string(cls, string: _str): ... @property def name(self) -> _str: ... def __hash__(self) -> int: ... def __eq__(self, other) -> bool: ... class PeriodDtype(PandasExtensionDtype): - type: type[Period] = ... - kind: _str = ... - str: _str = ... - base = ... - num: int = ... - def __new__(cls, freq=...): ... + def __new__(cls, freq: str | BaseOffset = ...): ... + def __hash__(self) -> int: ... + def __eq__(self, other) -> bool: ... @property def freq(self): ... - @classmethod - def construct_from_string(cls, string: _str): ... @property def name(self) -> _str: ... @property def na_value(self): ... - def __hash__(self) -> int: ... - def __eq__(self, other) -> bool: ... - @classmethod - def is_dtype(cls, dtype) -> bool: ... - @classmethod - def construct_array_type(cls): ... def __from_arrow__(self, array): ... class IntervalDtype(PandasExtensionDtype): - name: _str = ... - kind: _str = ... - str: _str = ... - base = ... - num: int = ... - def __new__(cls, subtype=...): ... + def __new__(cls, subtype: str | npt.DTypeLike | None = ...): ... + def __hash__(self) -> int: ... + def __eq__(self, other) -> bool: ... @property def subtype(self): ... - @classmethod - def construct_array_type(cls): ... - @classmethod - def construct_from_string(cls, string: _str): ... @property def type(self): ... - def __hash__(self) -> int: ... - def __eq__(self, other) -> bool: ... - @classmethod - def is_dtype(cls, dtype) -> bool: ... def __from_arrow__(self, array): ... diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index 9b06abfb6..3f2bb488b 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/pandas-stubs/core/indexes/timedeltas.pyi @@ -14,12 +14,23 @@ from pandas._libs import ( Timedelta, Timestamp, ) +from pandas._libs.tslibs import BaseOffset from pandas._typing import ( TimedeltaConvertibleTypes, num, ) class TimedeltaIndex(DatetimeTimedeltaMixin, TimedeltaIndexProperties): + def __init__( + self, + data=..., + unit: Literal["D", "h", "m", "s", "ms", "us", "ns"] = ..., + freq: str | BaseOffset = ..., + closed: object = ..., + dtype=..., + copy: bool = ..., + name: str = ..., + ): ... def __new__( cls, data=..., diff --git a/pyproject.toml b/pyproject.toml index f310eaa49..d8eae891d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,7 @@ tabulate = ">=0.8.10" jinja2 = "^3.1" scipy = ">=1.9.1" SQLAlchemy = "^1.4.41" +types-python-dateutil = ">=2.8.19" [build-system] requires = ["poetry-core>=1.0.0"] diff --git a/tests/test_arrays.py b/tests/test_arrays.py new file mode 100644 index 000000000..1afde80c8 --- /dev/null +++ b/tests/test_arrays.py @@ -0,0 +1,454 @@ +from typing import Type + +import numpy as np +import numpy.typing as npt +import pandas as pd +from pandas.arrays import ( + BooleanArray, + DatetimeArray, + IntegerArray, + IntervalArray, + PandasArray, + PeriodArray, + SparseArray, + StringArray, + TimedeltaArray, +) +from pandas.core.arrays.base import ExtensionArray +from typing_extensions import assert_type + +from pandas._libs.sparse import ( + BlockIndex, + IntIndex, + SparseIndex, +) + +from tests import check + +from pandas.tseries.offsets import Day + +LIST_MASK = [False, True, False, False, False, False, False, False, True, False] +ARRAY_MASK = np.array(LIST_MASK) + + +def test_integer_array() -> None: + ints = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) + check(assert_type(IntegerArray(ints, mask=ARRAY_MASK), IntegerArray), IntegerArray) + check( + assert_type(IntegerArray(ints, mask=ARRAY_MASK, copy=True), IntegerArray), + IntegerArray, + ) + + nulled_ints = [1, 2, 3, 4, 5, 6, 7, 8, None, 10] + check( + assert_type(pd.array(nulled_ints, dtype="UInt8"), Type[ExtensionArray]), + IntegerArray, + ) + check( + assert_type(pd.array(nulled_ints, dtype=pd.UInt8Dtype()), Type[ExtensionArray]), + IntegerArray, + ) + check( + assert_type(pd.array(nulled_ints, dtype=float), Type[ExtensionArray]), + PandasArray, + ) + check(assert_type(pd.array(ints, dtype=int), Type[ExtensionArray]), PandasArray) + + +def test_string_array() -> None: + strings = np.array(["a", "b", "c", "d", "e", "f", "g", "h", None, "j"]) + check(assert_type(StringArray(strings, copy=False), StringArray), StringArray) + check(assert_type(StringArray(strings, copy=True), StringArray), StringArray) + + strings_list = strings.tolist() + check( + assert_type(pd.array(strings_list, dtype="string"), Type[ExtensionArray]), + StringArray, + ) + check( + assert_type(pd.array(strings, dtype="string"), Type[ExtensionArray]), + StringArray, + ) + check(assert_type(pd.array(strings, dtype=str), Type[ExtensionArray]), PandasArray) + check(assert_type(pd.array(strings), Type[ExtensionArray]), StringArray) + + +def test_boolean_array() -> None: + bools = np.array([True, False, True, False, True, False, True, False, True, False]) + check(assert_type(BooleanArray(bools, mask=ARRAY_MASK), BooleanArray), BooleanArray) + check( + assert_type(BooleanArray(bools, mask=ARRAY_MASK, copy=True), BooleanArray), + BooleanArray, + ) + + nulled_bools = [True, False, True, False, True, False, True, False, None, False] + check(assert_type(pd.array(nulled_bools), Type[ExtensionArray]), BooleanArray) + check( + assert_type(pd.array(nulled_bools, dtype="bool"), Type[ExtensionArray]), + PandasArray, + ) + check( + assert_type(pd.array(nulled_bools, dtype=bool), Type[ExtensionArray]), + PandasArray, + ) + check( + assert_type( + pd.array(nulled_bools, dtype=pd.BooleanDtype()), Type[ExtensionArray] + ), + BooleanArray, + ) + + +def test_period_array() -> None: + p1 = pd.Period("2000-01-01", freq="D") + p2 = pd.Period("2000-01-02", freq="D") + pa = PeriodArray(pd.Series([p1, p2])) + check(assert_type(pa, PeriodArray), PeriodArray) + check(assert_type(PeriodArray(pd.Index([p1, p2])), PeriodArray), PeriodArray) + int_arr: npt.NDArray[np.int_] = np.ndarray([0, 1, 2]) + check(assert_type(PeriodArray(int_arr, freq="D"), PeriodArray), PeriodArray) + check( + assert_type(PeriodArray(np.ndarray([0, 1, 2]), freq=Day()), PeriodArray), + PeriodArray, + ) + check(assert_type(PeriodArray(pa), PeriodArray), PeriodArray) + dt = pd.PeriodDtype(freq="D") + period_idx = pd.Index([p1, p2]) + check( + assert_type(PeriodArray(period_idx, dtype=dt, copy=False), PeriodArray), + PeriodArray, + ) + + check( + assert_type( + PeriodArray(period_idx, dtype=dt, freq="D", copy=False), PeriodArray + ), + PeriodArray, + ) + + check(assert_type(pd.array([p1, p2]), Type[ExtensionArray]), PeriodArray) + check( + assert_type(pd.array([p1, p2], dtype="period[D]"), Type[ExtensionArray]), + PeriodArray, + ) + + +def test_datetime_array() -> None: + values = [pd.Timestamp("2000-1-1"), pd.Timestamp("2000-1-2")] + check( + assert_type( + DatetimeArray( + pd.Index(values), dtype=np.dtype("M8[ns]"), freq="D", copy=False + ), + DatetimeArray, + ), + DatetimeArray, + ) + check( + assert_type( + DatetimeArray( + pd.Series(values), dtype=np.dtype("M8[ns]"), freq="D", copy=False + ), + DatetimeArray, + ), + DatetimeArray, + ) + np_values = np.array([np.datetime64(1, "ns"), np.datetime64(2, "ns")]) + dta = DatetimeArray(np_values) + check(assert_type(DatetimeArray(dta), DatetimeArray), DatetimeArray) + data = np.array([1, 2, 3], dtype="M8[ns]") + check( + assert_type( + DatetimeArray(data, copy=False, dtype=pd.DatetimeTZDtype(tz="US/Central")), + DatetimeArray, + ), + DatetimeArray, + ) + + check(assert_type(pd.array(data), Type[ExtensionArray]), DatetimeArray) + check(assert_type(pd.array(np_values), Type[ExtensionArray]), DatetimeArray) + + +def test_interval_array_construction() -> None: + ia = IntervalArray([pd.Interval(0, 1), pd.Interval(1, 2)]) + check(assert_type(IntervalArray(ia), IntervalArray), IntervalArray) + check( + assert_type( + IntervalArray([pd.Interval(0, 1), pd.Interval(1, 2)], closed="left"), + IntervalArray, + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray([pd.Interval(0, 1), pd.Interval(1, 2)], closed="right"), + IntervalArray, + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray([pd.Interval(0, 1), pd.Interval(1, 2)], closed="both"), + IntervalArray, + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray([pd.Interval(0, 1), pd.Interval(1, 2)], closed="neither"), + IntervalArray, + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray( + [pd.Interval(0, 1), pd.Interval(1, 2)], + closed="neither", + verify_integrity=True, + ), + IntervalArray, + ), + IntervalArray, + ) + + check( + assert_type(IntervalArray.from_arrays([0, 1], [1, 2]), IntervalArray), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_arrays(np.array([0, 1]), np.array([1, 2])), IntervalArray + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_arrays(pd.Series([0, 1]), pd.Series([1, 2])), + IntervalArray, + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_arrays(pd.Index([0, 1]), pd.Index([1, 2])), IntervalArray + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_arrays([0, 1], [1, 2], closed="left", copy=False), + IntervalArray, + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_arrays( + [0, 1], [1, 2], closed="right", dtype=pd.IntervalDtype("int64") + ), + IntervalArray, + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_arrays( + [0, 1], [1, 2], closed="right", dtype=pd.IntervalDtype(float) + ), + IntervalArray, + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_arrays([0, 1], [1, 2], closed="both"), IntervalArray + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_arrays([0, 1], [1, 2], closed="neither"), IntervalArray + ), + IntervalArray, + ) + + breaks = [0, 1, 2, 3, 4.5] + check(assert_type(IntervalArray.from_breaks(breaks), IntervalArray), IntervalArray) + check( + assert_type( + IntervalArray.from_breaks(np.array(breaks), copy=False), IntervalArray + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_breaks(pd.Series(breaks), closed="left"), IntervalArray + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_breaks(pd.Index(breaks), closed="right"), IntervalArray + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_breaks(pd.Index(breaks), closed="both"), IntervalArray + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_breaks(pd.Index(breaks), closed="neither"), IntervalArray + ), + IntervalArray, + ) + check( + assert_type( + IntervalArray.from_breaks(pd.Index(breaks), dtype=pd.IntervalDtype(float)), + IntervalArray, + ), + IntervalArray, + ) + + +def test_integer_array_attrib_props() -> None: + ia = IntervalArray([pd.Interval(0, 1), pd.Interval(1, 2)]) + + ia.left + ia.right + ia.closed + ia.mid + ia.length + ia.is_empty + ia.is_non_overlapping_monotonic + + ia.contains(0.5) + ia.overlaps(pd.Interval(0.5, 1.0)) + ia.set_closed("right") + ia.set_closed("left") + ia.set_closed("both") + ia.set_closed("neither") + ia.to_tuples(True) + ia.to_tuples(False) + + +def test_timedelta_array() -> None: + td1, td2 = pd.Timedelta("1 days"), pd.Timedelta("2 days") + tda = TimedeltaArray(np.array([1, 2], dtype="timedelta64[ns]")) + check(assert_type(tda, TimedeltaArray), TimedeltaArray) + + tda = TimedeltaArray(np.array([1, 2], dtype="timedelta64[ns]"), copy=False) + tds = pd.Series([td1, td2]) + tdi = pd.Index([td1, td2]) + + check(assert_type(tda, TimedeltaArray), TimedeltaArray) + check(assert_type(TimedeltaArray(tds, freq="D"), TimedeltaArray), TimedeltaArray) + check(assert_type(TimedeltaArray(tds, freq=Day()), TimedeltaArray), TimedeltaArray) + check(assert_type(TimedeltaArray(tdi), TimedeltaArray), TimedeltaArray) + check(assert_type(TimedeltaArray(tda), TimedeltaArray), TimedeltaArray) + + check( + assert_type( + TimedeltaArray(tds, dtype=np.dtype("timedelta64[ns]")), TimedeltaArray + ), + TimedeltaArray, + ) + check( + assert_type( + TimedeltaArray(tds, dtype=np.dtype("timedelta64[ns]")), TimedeltaArray + ), + TimedeltaArray, + ) + + check( + assert_type( + pd.array(np.array([1, 2], dtype="timedelta64[ns]")), Type[ExtensionArray] + ), + TimedeltaArray, + ) + check(assert_type(pd.array(tdi), Type[ExtensionArray]), TimedeltaArray) + check(assert_type(pd.array(tds, copy=False), Type[ExtensionArray]), TimedeltaArray) + + +def test_sparse_array() -> None: + ints = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0] + nulled_ints = [1, 2, 3, 4, 5, 6, 7, 8, 9, np.nan] + zero_ints = [0, 2, 3, 4, 5, 6, 0, 8, 9, 0] + + check(assert_type(SparseArray(3.0, fill_value=np.nan), SparseArray), SparseArray) + check( + assert_type( + SparseArray(nulled_ints, fill_value=np.nan), + SparseArray, + ), + SparseArray, + ) + sa = SparseArray(zero_ints, fill_value=0) + check(assert_type(sa.sp_index, SparseIndex), IntIndex) + check( + assert_type(SparseArray(sa.sp_values, sparse_index=sa.sp_index), SparseArray), + SparseArray, + ) + sa_block = SparseArray(zero_ints, fill_value=0, kind="block") + check(assert_type(sa_block.sp_index, SparseIndex), BlockIndex) + check( + assert_type( + SparseArray(sa_block.sp_values, sparse_index=sa_block.sp_index), SparseArray + ), + SparseArray, + ) + + check( + assert_type( + SparseArray( + [True, False, False, False, False, False, False, True, False, False], + fill_value=False, + ), + SparseArray, + ), + SparseArray, + ) + check( + assert_type( + SparseArray( + [ + pd.Timestamp("2011-01-01"), + pd.Timestamp("2011-01-02"), + pd.Timestamp("2011-01-03"), + pd.NaT, + pd.NaT, + ], + fill_value=pd.NaT, + ), + SparseArray, + ), + SparseArray, + ) + check( + assert_type( + SparseArray([pd.Timedelta(days=1), pd.NaT, pd.NaT], fill_value=pd.NaT), + SparseArray, + ), + SparseArray, + ) + + check( + assert_type( + SparseArray(nulled_ints, kind="integer", copy=False), + SparseArray, + ), + SparseArray, + ) + check( + assert_type( + SparseArray(nulled_ints, kind="block", copy=True), + SparseArray, + ), + SparseArray, + ) + check(assert_type(SparseArray(ints, dtype="i4"), SparseArray), SparseArray) + check(assert_type(SparseArray(ints, dtype="int32"), SparseArray), SparseArray) + check(assert_type(SparseArray(ints, dtype=np.int16), SparseArray), SparseArray) diff --git a/tests/test_dtypes.py b/tests/test_dtypes.py new file mode 100644 index 000000000..8db3e3ff1 --- /dev/null +++ b/tests/test_dtypes.py @@ -0,0 +1,118 @@ +from datetime import ( + timedelta, + timezone, +) + +import numpy as np +import pandas as pd +import pyarrow as pa +from typing_extensions import assert_type + +from tests import check + +from pandas.tseries.offsets import ( + BusinessDay, + CustomBusinessDay, + Day, +) + + +def test_datetimetz_dtype() -> None: + check( + assert_type(pd.DatetimeTZDtype(unit="ns", tz="UTC"), pd.DatetimeTZDtype), + pd.DatetimeTZDtype, + ) + check( + assert_type( + pd.DatetimeTZDtype(unit="ns", tz=timezone(timedelta(hours=1))), + pd.DatetimeTZDtype, + ), + pd.DatetimeTZDtype, + ) + + +def test_period_dtype() -> None: + check(assert_type(pd.PeriodDtype(freq="D"), pd.PeriodDtype), pd.PeriodDtype) + check(assert_type(pd.PeriodDtype(freq=Day()), pd.PeriodDtype), pd.PeriodDtype) + check( + assert_type(pd.PeriodDtype(freq=BusinessDay()), pd.PeriodDtype), pd.PeriodDtype + ) + check( + assert_type(pd.PeriodDtype(freq=CustomBusinessDay()), pd.PeriodDtype), + pd.PeriodDtype, + ) + + +def test_interval_dtype() -> None: + check( + assert_type( + pd.Interval(pd.Timestamp("2017-01-01"), pd.Timestamp("2017-01-02")), + "pd.Interval[pd.Timestamp]", + ), + pd.Interval, + ) + check( + assert_type(pd.Interval(1, 2, closed="left"), "pd.Interval[int]"), pd.Interval + ) + check( + assert_type(pd.Interval(1.0, 2.5, closed="right"), "pd.Interval[float]"), + pd.Interval, + ) + check( + assert_type(pd.Interval(1.0, 2.5, closed="both"), "pd.Interval[float]"), + pd.Interval, + ) + check( + assert_type( + pd.Interval( + pd.Timedelta("1 day"), pd.Timedelta("2 days"), closed="neither" + ), + "pd.Interval[pd.Timedelta]", + ), + pd.Interval, + ) + + +def test_int64_dtype() -> None: + check(assert_type(pd.Int64Dtype(), pd.Int64Dtype), pd.Int64Dtype) + + +def test_categorical_dtype() -> None: + check( + assert_type( + pd.CategoricalDtype(categories=["a", "b", "c"], ordered=True), + pd.CategoricalDtype, + ), + pd.CategoricalDtype, + ) + check( + assert_type(pd.CategoricalDtype(categories=[1, 2, 3]), pd.CategoricalDtype), + pd.CategoricalDtype, + ) + + +def test_sparse_dtype() -> None: + check(assert_type(pd.SparseDtype(str), pd.SparseDtype), pd.SparseDtype) + check(assert_type(pd.SparseDtype(complex), pd.SparseDtype), pd.SparseDtype) + check(assert_type(pd.SparseDtype(bool), pd.SparseDtype), pd.SparseDtype) + check(assert_type(pd.SparseDtype(int), pd.SparseDtype), pd.SparseDtype) + check(assert_type(pd.SparseDtype(np.int64), pd.SparseDtype), pd.SparseDtype) + check(assert_type(pd.SparseDtype(str), pd.SparseDtype), pd.SparseDtype) + check(assert_type(pd.SparseDtype(float), pd.SparseDtype), pd.SparseDtype) + check(assert_type(pd.SparseDtype(np.datetime64), pd.SparseDtype), pd.SparseDtype) + check(assert_type(pd.SparseDtype(np.timedelta64), pd.SparseDtype), pd.SparseDtype) + check(assert_type(pd.SparseDtype("datetime64"), pd.SparseDtype), pd.SparseDtype) + check(assert_type(pd.SparseDtype(), pd.SparseDtype), pd.SparseDtype) + + +def test_string_dtype() -> None: + check(assert_type(pd.StringDtype("pyarrow"), pd.StringDtype), pd.StringDtype) + check(assert_type(pd.StringDtype("python"), pd.StringDtype), pd.StringDtype) + + +def test_boolean_dtype() -> None: + check(assert_type(pd.BooleanDtype(), pd.BooleanDtype), pd.BooleanDtype) + + +def test_arrow_dtype() -> None: + check(assert_type(pd.ArrowDtype(pa.int64()), pd.ArrowDtype), pd.ArrowDtype) diff --git a/tests/test_scalars.py b/tests/test_scalars.py new file mode 100644 index 000000000..aa2c9f1ac --- /dev/null +++ b/tests/test_scalars.py @@ -0,0 +1,1097 @@ +from __future__ import annotations + +import datetime as dt +from typing import ( + TYPE_CHECKING, + Any, + Literal, + Optional, + cast, +) + +import dateutil.tz +import numpy as np +from numpy import typing as npt +import pandas as pd +import pytest +import pytz +from typing_extensions import assert_type + +from pandas._libs.tslibs import BaseOffset +from pandas._libs.tslibs.timedeltas import Components + +if TYPE_CHECKING: + from pandas.core.series import PeriodSeries # noqa: F401 + from pandas.core.series import TimedeltaSeries # noqa: F401 + from pandas.core.series import TimestampSeries # noqa: F401 + + from pandas._typing import np_ndarray_bool +else: + np_ndarray_bool = Any + +from tests import check + +from pandas.tseries.offsets import Day + + +def test_period() -> None: + p = pd.Period("2012-1-1", freq="D") + check(assert_type(p, pd.Period), pd.Period) + check(assert_type(pd.Period(p), pd.Period), pd.Period) + check(assert_type(pd.Period("2012-1-1", freq=Day()), pd.Period), pd.Period) + check( + assert_type(pd.Period(freq="D", year=2012, day=1, month=1), pd.Period), + pd.Period, + ) + check( + assert_type(pd.Period(None, "D", year=2012, day=1, month=1), pd.Period), + pd.Period, + ) + check( + assert_type(pd.Period(None, "D", 1, year=2012, day=1, month=1), pd.Period), + pd.Period, + ) + check( + assert_type( + pd.Period( + freq="s", year=2012, month=1, day=1, hour=12, minute=30, second=45 + ), + pd.Period, + ), + pd.Period, + ) + check(assert_type(pd.Period(freq="Q", year=2012, quarter=2), pd.Period), pd.Period) + check(assert_type(p.day, int), int) + check(assert_type(p.day_of_week, int), int) + check(assert_type(p.day_of_year, int), int) + check(assert_type(p.dayofweek, int), int) + check(assert_type(p.dayofyear, int), int) + check(assert_type(p.days_in_month, int), int) + check(assert_type(p.daysinmonth, int), int) + check(assert_type(p.end_time, pd.Timestamp), pd.Timestamp) + check(assert_type(p.freqstr, str), str) + check(assert_type(p.hour, int), int) + check(assert_type(p.is_leap_year, bool), bool) + check(assert_type(p.minute, int), int) + check(assert_type(p.month, int), int) + check(assert_type(p.quarter, int), int) + check(assert_type(p.qyear, int), int) + check(assert_type(p.second, int), int) + check(assert_type(p.start_time, pd.Timestamp), pd.Timestamp) + check(assert_type(p.week, int), int) + check(assert_type(p.weekday, int), int) + check(assert_type(p.weekofyear, int), int) + check(assert_type(p.year, int), int) + check(assert_type(p.freq, BaseOffset), Day) + check(assert_type(p.ordinal, int), int) + + p2 = pd.Period("2012-1-1", freq="2D") + check(assert_type(p2.freq, BaseOffset), Day) + + as0 = pd.Timedelta(1, "D") + as1 = dt.timedelta(days=1) + as2 = np.timedelta64(1, "D") + as3 = np.int64(1) + as4 = int(1) + as5 = pd.period_range("2012-1-1", periods=10, freq="D") + as6 = pd.Period("2012-1-1", freq="D") + as7 = cast("TimedeltaSeries", pd.Series([pd.Timedelta(days=1)])) + as8 = cast("PeriodSeries", pd.Series([as6])) + + check(assert_type(p + as0, pd.Period), pd.Period) + check(assert_type(p + as1, pd.Period), pd.Period) + check(assert_type(p + as2, pd.Period), pd.Period) + check(assert_type(p + as3, pd.Period), pd.Period) + check(assert_type(p + as4, pd.Period), pd.Period) + check(assert_type(p + p.freq, pd.Period), pd.Period) + check(assert_type(p + (p - as5), pd.PeriodIndex), pd.PeriodIndex) + check(assert_type(p + as7, "PeriodSeries"), pd.Series) + das8 = cast("TimedeltaSeries", (as8 - as8)) + check(assert_type(p + das8, "PeriodSeries"), pd.Series) + check(assert_type(p - as0, pd.Period), pd.Period) + check(assert_type(p - as1, pd.Period), pd.Period) + check(assert_type(p - as2, pd.Period), pd.Period) + check(assert_type(p - as3, pd.Period), pd.Period) + check(assert_type(p - as4, pd.Period), pd.Period) + check(assert_type(p - as5, pd.Index), pd.Index) + check(assert_type(p - as6, BaseOffset), Day) + check(assert_type(p - as7, "PeriodSeries"), pd.Series) + check(assert_type(p - p.freq, pd.Period), pd.Period) + + check(assert_type(as0 + p, pd.Period), pd.Period) + check(assert_type(as1 + p, pd.Period), pd.Period) + check(assert_type(as2 + p, pd.Period), pd.Period) + check(assert_type(as3 + p, pd.Period), pd.Period) + check(assert_type(as4 + p, pd.Period), pd.Period) + check(assert_type(as7 + p, "PeriodSeries"), pd.Series) + check(assert_type(p.freq + p, pd.Period), pd.Period) + # TODO: PeriodIndex should have a __sub__ with correct types, this op is valid + # and so the assert_type is skipped + check(as5 - p, pd.Index) # type: ignore[operator] + + check(assert_type(p.__radd__(as0), pd.Period), pd.Period) + check(assert_type(p.__radd__(as1), pd.Period), pd.Period) + check(assert_type(p.__radd__(as2), pd.Period), pd.Period) + check(assert_type(p.__radd__(as3), pd.Period), pd.Period) + check(assert_type(p.__radd__(as4), pd.Period), pd.Period) + check(assert_type(p.__radd__(p.freq), pd.Period), pd.Period) + + c0 = pd.Period("2012-1-1", freq="D") + c1 = pd.period_range("2012-1-1", periods=10, freq="D") + + check(assert_type(p == c0, bool), bool) + check(assert_type(p == c1, np_ndarray_bool), np.ndarray) + check(assert_type(c0 == p, bool), bool) + check(assert_type(c1 == p, np_ndarray_bool), np.ndarray) + + check(assert_type(p != c0, bool), bool) + check(assert_type(p != c1, np_ndarray_bool), np.ndarray) + check(assert_type(c0 != p, bool), bool) + check(assert_type(c1 != p, np_ndarray_bool), np.ndarray) + + check(assert_type(p > c0, bool), bool) + check(assert_type(p > c1, np_ndarray_bool), np.ndarray) + check(assert_type(c0 > p, bool), bool) + check(assert_type(c1 > p, np_ndarray_bool), np.ndarray) + + check(assert_type(p < c0, bool), bool) + check(assert_type(p < c1, np_ndarray_bool), np.ndarray) + check(assert_type(c0 < p, bool), bool) + check(assert_type(c1 < p, np_ndarray_bool), np.ndarray) + + check(assert_type(p <= c0, bool), bool) + check(assert_type(p <= c1, np_ndarray_bool), np.ndarray) + check(assert_type(c0 <= p, bool), bool) + check(assert_type(c1 <= p, np_ndarray_bool), np.ndarray) + + check(assert_type(p >= c0, bool), bool) + check(assert_type(p >= c1, np_ndarray_bool), np.ndarray) + check(assert_type(c0 >= p, bool), bool) + check(assert_type(c1 >= p, np_ndarray_bool), np.ndarray) + + p3 = pd.Period("2007-01", freq="M") + check(assert_type(p3.to_timestamp("D", "S"), pd.Timestamp), pd.Timestamp) + check(assert_type(p3.to_timestamp("D", "E"), pd.Timestamp), pd.Timestamp) + check(assert_type(p3.to_timestamp("D", "Finish"), pd.Timestamp), pd.Timestamp) + check(assert_type(p3.to_timestamp("D", "End"), pd.Timestamp), pd.Timestamp) + check(assert_type(p3.to_timestamp("D", "Begin"), pd.Timestamp), pd.Timestamp) + check(assert_type(p3.to_timestamp("D", "Start"), pd.Timestamp), pd.Timestamp) + + check(assert_type(p3.asfreq("D", "S"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "E"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "Finish"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "Begin"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "Start"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "End"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "end"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "start"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "begin"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "finish"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "s"), pd.Period), pd.Period) + check(assert_type(p3.asfreq(Day(), "e"), pd.Period), pd.Period) + + check(assert_type(pd.Period.now("D"), pd.Period), pd.Period) + check(assert_type(pd.Period.now(Day()), pd.Period), pd.Period) + + check(assert_type(p.strftime("%Y-%m-%d"), str), str) + check(assert_type(hash(p), int), int) + + +def test_timedelta() -> None: + check(assert_type(pd.Timedelta(1, "W"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "w"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "D"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "d"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "days"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "day"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "hours"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "hour"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "hr"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "h"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "m"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "minute"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "min"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "minutes"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "t"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "s"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "seconds"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "sec"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "second"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "ms"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "milliseconds"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "millisecond"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "milli"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "millis"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "l"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "us"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "microseconds"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "microsecond"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "µs"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "micro"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "micros"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "u"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "ns"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "nanoseconds"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "nano"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "nanos"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "nanosecond"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(1, "n"), pd.Timedelta), pd.Timedelta) + + check(assert_type(pd.Timedelta("1 W"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 w"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 D"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 d"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 days"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 day"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 hours"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 hour"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 hr"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 h"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 m"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 minute"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 min"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 minutes"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 t"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 s"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 seconds"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 sec"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 second"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 ms"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 milliseconds"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 millisecond"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 milli"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 millis"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 l"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 us"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 microseconds"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 microsecond"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 µs"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 micro"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 micros"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 u"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 ns"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 nanoseconds"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 nano"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 nanos"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 nanosecond"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta("1 n"), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(days=1), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(seconds=1), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(microseconds=1), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(minutes=1), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(hours=1), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(weeks=1), pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Timedelta(milliseconds=1), pd.Timedelta), pd.Timedelta) + + td = pd.Timedelta("1 day") + check(assert_type(td.value, int), int) + check(assert_type(td.asm8, np.timedelta64), np.timedelta64) + + check(assert_type(td.days, int), int) + check(assert_type(td.microseconds, int), int) + check(assert_type(td.nanoseconds, int), int) + check(assert_type(td.seconds, int), int) + check(assert_type(td.value, int), int) + check(assert_type(td.resolution_string, str), str) + check(assert_type(td.components, Components), Components) + + check(assert_type(td.ceil("D"), pd.Timedelta), pd.Timedelta) + check(assert_type(td.floor(Day()), pd.Timedelta), pd.Timedelta) + check(assert_type(td.isoformat(), str), str) + check(assert_type(td.round("s"), pd.Timedelta), pd.Timedelta) + check(assert_type(td.to_numpy(), np.timedelta64), np.timedelta64) + check(assert_type(td.to_pytimedelta(), dt.timedelta), dt.timedelta) + check(assert_type(td.to_timedelta64(), np.timedelta64), np.timedelta64) + check(assert_type(td.total_seconds(), float), float) + check(assert_type(td.view(np.int64), object), np.int64) + check(assert_type(td.view("i8"), object), np.int64) + + ndarray_td64: npt.NDArray[np.timedelta64] = np.array( + [1, 2, 3], dtype="timedelta64[D]" + ) + ndarray_dt64: npt.NDArray[np.datetime64] = np.array( + [1, 2, 3], dtype="datetime64[D]" + ) + + check(assert_type(td + pd.Period("2012-01-01", freq="D"), pd.Period), pd.Period) + check(assert_type(td + pd.Timestamp("2012-01-01"), pd.Timestamp), pd.Timestamp) + check(assert_type(td + dt.datetime(2012, 1, 1), pd.Timestamp), pd.Timestamp) + check(assert_type(td + dt.date(2012, 1, 1), dt.date), dt.date) + check(assert_type(td + np.datetime64(1, "ns"), pd.Timestamp), pd.Timestamp) + check(assert_type(td + dt.timedelta(days=1), pd.Timedelta), pd.Timedelta) + check(assert_type(td + np.timedelta64(1, "D"), pd.Timedelta), pd.Timedelta) + check( + assert_type( + td + pd.period_range("2012-01-01", periods=3, freq="D"), pd.PeriodIndex + ), + pd.PeriodIndex, + ) + check( + assert_type(td + pd.date_range("2012-01-01", periods=3), pd.DatetimeIndex), + pd.DatetimeIndex, + ) + check( + assert_type( + td + ndarray_td64, + npt.NDArray[np.timedelta64], + ), + np.ndarray, + ) + check( + assert_type( + td + ndarray_dt64, # pyright: ignore[reportGeneralTypeIssues] + npt.NDArray[np.datetime64], + ), + np.ndarray, + ) + + check(assert_type(td - td, pd.Timedelta), pd.Timedelta) + check(assert_type(td - dt.timedelta(days=1), pd.Timedelta), pd.Timedelta) + check(assert_type(td - np.timedelta64(1, "D"), pd.Timedelta), pd.Timedelta) + check( + assert_type( + td - ndarray_td64, + npt.NDArray[np.timedelta64], + ), + np.ndarray, + ) + # pyright appears to get some things wrong when __rsub__ is called, + # hence pyright ignores + check(assert_type(pd.Period("2012-01-01", freq="D") - td, pd.Period), pd.Period) + check(assert_type(pd.Timestamp("2012-01-01") - td, pd.Timestamp), pd.Timestamp) + check(assert_type(dt.datetime(2012, 1, 1) - td, dt.datetime), dt.datetime) + check(assert_type(dt.date(2012, 1, 1) - td, dt.date), dt.date) + check(assert_type(np.datetime64(1, "ns") - td, pd.Timestamp), pd.Timestamp) + check( + assert_type( + dt.timedelta(days=1) - td, # pyright: ignore[reportGeneralTypeIssues] + pd.Timedelta, + ), + pd.Timedelta, + ) + check(assert_type(np.timedelta64(1, "D") - td, pd.Timedelta), pd.Timedelta) + check( + assert_type( + pd.period_range("2012-01-01", periods=3, freq="D") - td, pd.PeriodIndex + ), + pd.PeriodIndex, + ) + check( + assert_type(pd.date_range("2012-01-01", periods=3) - td, pd.DatetimeIndex), + pd.DatetimeIndex, + ) + check( + assert_type( + ndarray_td64 - td, # pyright: ignore[reportGeneralTypeIssues] + npt.NDArray[np.timedelta64], + ), + np.ndarray, + ) + check( + assert_type( + ndarray_dt64 - td, # pyright: ignore[reportGeneralTypeIssues] + npt.NDArray[np.datetime64], + ), + np.ndarray, + ) + + with pytest.warns(FutureWarning): + i_idx = cast(pd.Int64Index, pd.Index([1, 2, 3], dtype=int)) + f_idx = cast(pd.Float64Index, pd.Index([1.2, 2.2, 3.4], dtype=float)) + + check(assert_type(td * 3, pd.Timedelta), pd.Timedelta) + check(assert_type(td * 3.5, pd.Timedelta), pd.Timedelta) + check(assert_type(td * np.array([1, 2, 3]), np.ndarray), np.ndarray) + check(assert_type(td * np.array([1.2, 2.2, 3.4]), np.ndarray), np.ndarray) + check(assert_type(td * pd.Series([1, 2, 3]), pd.Series), pd.Series) + check(assert_type(td * pd.Series([1.2, 2.2, 3.4]), pd.Series), pd.Series) + check(assert_type(td * i_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td * f_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) + + check(assert_type(3 * td, pd.Timedelta), pd.Timedelta) + check(assert_type(3.5 * td, pd.Timedelta), pd.Timedelta) + check(assert_type(pd.Series([1, 2, 3]) * td, "TimedeltaSeries"), pd.Series) + check(assert_type(pd.Series([1.2, 2.2, 3.4]) * td, "TimedeltaSeries"), pd.Series) + check(assert_type(i_idx * td, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(f_idx * td, pd.TimedeltaIndex), pd.TimedeltaIndex) + + np_intp_arr: npt.NDArray[np.integer] = np.array([1, 2, 3]) + np_float_arr: npt.NDArray[np.floating] = np.array([1, 2, 3]) + check(assert_type(td // td, int), int) + check(assert_type(td // 3, pd.Timedelta), pd.Timedelta) + check(assert_type(td // 3.5, pd.Timedelta), pd.Timedelta) + check(assert_type(td // np_intp_arr, npt.NDArray[np.timedelta64]), np.ndarray) + check(assert_type(td // np_float_arr, npt.NDArray[np.timedelta64]), np.ndarray) + check(assert_type(td // pd.Series([1, 2, 3]), pd.Series), pd.Series) + check(assert_type(td // pd.Series([1.2, 2.2, 3.4]), pd.Series), pd.Series) + check(assert_type(td // i_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td // f_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) + + # Note: None of the reverse floordiv work + # check(assert_type(3 // td, pd.Timedelta), pd.Timedelta) + # check(assert_type(3.5// td, pd.Timedelta), pd.Timedelta) + # check(assert_type(np_intp_arr// td, npt.NDArray[np.timedelta64]), np.ndarray) + # check(assert_type(np_float_arr// td, npt.NDArray[np.timedelta64]), np.ndarray) + # check(assert_type(pd.Series([1, 2, 3])// td, pd.Series), pd.Series) + # check(assert_type(pd.Series([1.2, 2.2, 3.4])// td, pd.Series), pd.Series) + # check(assert_type(i_idx, pd.TimedeltaIndex)// td, pd.TimedeltaIndex) + # check( assert_type(f_idx// td, pd.TimedeltaIndex), pd.TimedeltaIndex) + + check(assert_type(td / td, float), float) + check(assert_type(td / 3, pd.Timedelta), pd.Timedelta) + check(assert_type(td / 3.5, pd.Timedelta), pd.Timedelta) + check( + assert_type(td / np.array([1, 2, 3]), npt.NDArray[np.timedelta64]), np.ndarray + ) + check( + assert_type(td / np.array([1.2, 2.2, 3.4]), npt.NDArray[np.timedelta64]), + np.ndarray, + ) + check(assert_type(td / pd.Series([1, 2, 3]), pd.Series), pd.Series) + check(assert_type(td / pd.Series([1.2, 2.2, 3.4]), pd.Series), pd.Series) + check(assert_type(td / i_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td / f_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) + + # Note: None of the reverse truediv work + # check(assert_type(3 / td, pd.Timedelta), pd.Timedelta) + # check(assert_type(3.5 / td, pd.Timedelta), pd.Timedelta) + # check(assert_type(np.array([1, 2, 3]) / td, npt.NDArray[np.timedelta64]), np.ndarray) + # check(assert_type(np.array([1.2, 2.2, 3.4]) / td, npt.NDArray[np.timedelta64]),np.ndarray,) + # check(assert_type(pd.Series([1, 2, 3]) / td, pd.Series), pd.Series) + # check(assert_type(pd.Series([1.2, 2.2, 3.4]) / td, pd.Series), pd.Series) + # check(assert_type(i_idx / td, pd.TimedeltaIndex), pd.TimedeltaIndex) + # check(assert_type(f_idx / td, pd.TimedeltaIndex), pd.TimedeltaIndex) + + check(assert_type(td % 3, pd.Timedelta), pd.Timedelta) + check(assert_type(td % 3.5, pd.Timedelta), pd.Timedelta) + check(assert_type(td % td, pd.Timedelta), pd.Timedelta) + check( + assert_type(td % np.array([1, 2, 3]), npt.NDArray[np.timedelta64]), np.ndarray + ) + check( + assert_type(td % np.array([1.2, 2.2, 3.4]), npt.NDArray[np.timedelta64]), + np.ndarray, + ) + check(assert_type(td % pd.Series([1, 2, 3]), pd.Series), pd.Series) + check(assert_type(td % pd.Series([1.2, 2.2, 3.4]), pd.Series), pd.Series) + check(assert_type(td % i_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) + + check( + assert_type(td % f_idx, pd.TimedeltaIndex), + pd.TimedeltaIndex, + ) + + check(assert_type(td.__abs__(), pd.Timedelta), pd.Timedelta) + check(assert_type(-td, pd.Timedelta), pd.Timedelta) + check(assert_type(+td, pd.Timedelta), pd.Timedelta) + + check(assert_type(td < td, bool), bool) + check(assert_type(td < dt.timedelta(days=1), bool), bool) + check(assert_type(td < np.timedelta64(1, "D"), bool), bool) + check(assert_type(td < ndarray_td64, np_ndarray_bool), np.ndarray) + check( + assert_type(pd.TimedeltaIndex([1, 2, 3], unit="D") < td, np_ndarray_bool), + np.ndarray, + ) + check(assert_type(dt.timedelta(days=1) < td, bool), bool) + check(assert_type(ndarray_td64 < td, np_ndarray_bool), np.ndarray) + check( + assert_type(pd.TimedeltaIndex([1, 2, 3], unit="D") < td, np_ndarray_bool), + np.ndarray, + ) + + check(assert_type(td > td, bool), bool) + check(assert_type(td > dt.timedelta(days=1), bool), bool) + check(assert_type(td > np.timedelta64(1, "D"), bool), bool) + check(assert_type(td > ndarray_td64, np_ndarray_bool), np.ndarray) + check( + assert_type(td > pd.TimedeltaIndex([1, 2, 3], unit="D"), np_ndarray_bool), + np.ndarray, + ) + check(assert_type(dt.timedelta(days=1) > td, bool), bool) + check(assert_type(ndarray_td64 > td, np_ndarray_bool), np.ndarray) + check( + assert_type(pd.TimedeltaIndex([1, 2, 3], unit="D") > td, np_ndarray_bool), + np.ndarray, + ) + + check(assert_type(td <= td, bool), bool) + check(assert_type(td <= dt.timedelta(days=1), bool), bool) + check(assert_type(td <= np.timedelta64(1, "D"), bool), bool) + check(assert_type(td <= ndarray_td64, np_ndarray_bool), np.ndarray) + check( + assert_type(td <= pd.TimedeltaIndex([1, 2, 3], unit="D"), np_ndarray_bool), + np.ndarray, + ) + check(assert_type(dt.timedelta(days=1) <= td, bool), bool) + check(assert_type(ndarray_td64 <= td, np_ndarray_bool), np.ndarray) + check( + assert_type(pd.TimedeltaIndex([1, 2, 3], unit="D") <= td, np_ndarray_bool), + np.ndarray, + ) + + check(assert_type(td >= td, bool), bool) + check(assert_type(td >= dt.timedelta(days=1), bool), bool) + check(assert_type(td >= np.timedelta64(1, "D"), bool), bool) + check(assert_type(td >= ndarray_td64, np_ndarray_bool), np.ndarray) + check( + assert_type(td >= pd.TimedeltaIndex([1, 2, 3], unit="D"), np_ndarray_bool), + np.ndarray, + ) + check(assert_type(dt.timedelta(days=1) >= td, bool), bool) + check(assert_type(ndarray_td64 >= td, np_ndarray_bool), np.ndarray) + check( + assert_type(pd.TimedeltaIndex([1, 2, 3], unit="D") >= td, np_ndarray_bool), + np.ndarray, + ) + + check(assert_type(td == td, bool), bool) + check(assert_type(td == dt.timedelta(days=1), bool), bool) + check(assert_type(td == np.timedelta64(1, "D"), bool), bool) + check(assert_type(td == ndarray_td64, np_ndarray_bool), np.ndarray) + check( + assert_type(td == pd.TimedeltaIndex([1, 2, 3], unit="D"), np_ndarray_bool), + np.ndarray, + ) + check(assert_type(td == pd.Series([1, 2, 3]), pd.Series), pd.Series) + check(assert_type(dt.timedelta(days=1) == td, bool), bool) + check( + assert_type(pd.TimedeltaIndex([1, 2, 3], unit="D") == td, np_ndarray_bool), + np.ndarray, + ) + check(assert_type(pd.Series([1, 2, 3]) == td, "pd.Series[bool]"), pd.Series) + + check(assert_type(td == 1, bool), bool) + check(assert_type(td == (3 + 2j), bool), bool) + + check(assert_type(td != td, bool), bool) + check(assert_type(td != dt.timedelta(days=1), bool), bool) + check(assert_type(td != np.timedelta64(1, "D"), bool), bool) + check(assert_type(td != ndarray_td64, np_ndarray_bool), np.ndarray) + check( + assert_type(td != pd.TimedeltaIndex([1, 2, 3], unit="D"), np_ndarray_bool), + np.ndarray, + ) + check(assert_type(td != pd.Series([1, 2, 3]), pd.Series), pd.Series) + check(assert_type(dt.timedelta(days=1) != td, bool), bool) + check( + assert_type(pd.TimedeltaIndex([1, 2, 3], unit="D") != td, np_ndarray_bool), + np.ndarray, + ) + check(assert_type(pd.Series([1, 2, 3]) != td, "pd.Series[bool]"), pd.Series) + check(assert_type(td != 1, bool), bool) + check(assert_type(td != (3 + 2j), bool), bool) + + # Mismatch due to NumPy ops returning Any + check(assert_type(np.array([1, 2, 3]) * td, Any), np.ndarray) + check(assert_type(np.array([1.2, 2.2, 3.4]) * td, Any), np.ndarray) + check(assert_type(np.timedelta64(1, "D") < td, Any), np.bool_) + check(assert_type(np.timedelta64(1, "D") > td, Any), np.bool_) + check(assert_type(np.timedelta64(1, "D") <= td, Any), np.bool_) + check(assert_type(np.timedelta64(1, "D") >= td, Any), np.bool_) + check(assert_type(np.timedelta64(1, "D") == td, Any), np.bool_) + check(assert_type(ndarray_td64 == td, Any), np.ndarray) + check(assert_type(ndarray_td64 != td, Any), np.ndarray) + check(assert_type(np.timedelta64(1, "D") != td, Any), np.bool_) + + +def test_interval() -> None: + i0 = pd.Interval(0, 1, closed="left") + i1 = pd.Interval(0.0, 1.0, closed="right") + i2 = pd.Interval( + pd.Timestamp("2017-01-01"), pd.Timestamp("2017-01-02"), closed="both" + ) + i3 = pd.Interval(pd.Timedelta("1 days"), pd.Timedelta("2 days"), closed="neither") + check(assert_type(i0, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1, "pd.Interval[float]"), pd.Interval) + check(assert_type(i2, "pd.Interval[pd.Timestamp]"), pd.Interval) + check(assert_type(i3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0.closed, Literal["left", "right", "both", "neither"]), str) + check(assert_type(i0.closed_left, bool), bool) + check(assert_type(i0.closed_right, bool), bool) + check(assert_type(i0.is_empty, bool), bool) + check(assert_type(i0.left, int), int) + check(assert_type(i0.length, int), int) + check(assert_type(i0.mid, float), float) + check(assert_type(i0.open_left, bool), bool) + check(assert_type(i0.open_right, bool), bool) + check(assert_type(i0.right, int), int) + + check(assert_type(i1.closed, Literal["left", "right", "both", "neither"]), str) + check(assert_type(i1.closed_left, bool), bool) + check(assert_type(i1.closed_right, bool), bool) + check(assert_type(i1.is_empty, bool), bool) + check(assert_type(i1.left, float), float) + check(assert_type(i1.length, float), float) + check(assert_type(i1.mid, float), float) + check(assert_type(i1.open_left, bool), bool) + check(assert_type(i1.open_right, bool), bool) + check(assert_type(i1.right, float), float) + + check(assert_type(i2.closed, Literal["left", "right", "both", "neither"]), str) + check(assert_type(i2.closed_left, bool), bool) + check(assert_type(i2.closed_right, bool), bool) + check(assert_type(i2.is_empty, bool), bool) + check(assert_type(i2.left, pd.Timestamp), pd.Timestamp) + check(assert_type(i2.length, pd.Timedelta), pd.Timedelta) + check(assert_type(i2.mid, pd.Timestamp), pd.Timestamp) + check(assert_type(i2.open_left, bool), bool) + check(assert_type(i2.open_right, bool), bool) + check(assert_type(i2.right, pd.Timestamp), pd.Timestamp) + + check(assert_type(i3.closed, Literal["left", "right", "both", "neither"]), str) + check(assert_type(i3.closed_left, bool), bool) + check(assert_type(i3.closed_right, bool), bool) + check(assert_type(i3.is_empty, bool), bool) + check(assert_type(i3.left, pd.Timedelta), pd.Timedelta) + check(assert_type(i3.length, pd.Timedelta), pd.Timedelta) + check(assert_type(i3.mid, pd.Timedelta), pd.Timedelta) + check(assert_type(i3.open_left, bool), bool) + check(assert_type(i3.open_right, bool), bool) + check(assert_type(i3.right, pd.Timedelta), pd.Timedelta) + + check(assert_type(i0.overlaps(pd.Interval(0.5, 1.5, closed="left")), bool), bool) + check(assert_type(i0.overlaps(pd.Interval(2, 3, closed="left")), bool), bool) + + check(assert_type(i1.overlaps(pd.Interval(0.5, 1.5, closed="left")), bool), bool) + check(assert_type(i1.overlaps(pd.Interval(2, 3, closed="left")), bool), bool) + ts1 = pd.Timestamp(year=2017, month=1, day=1) + ts2 = pd.Timestamp(year=2017, month=1, day=2) + check(assert_type(i2.overlaps(pd.Interval(ts1, ts2, closed="left")), bool), bool) + td1 = pd.Timedelta(days=1) + td2 = pd.Timedelta(days=3) + check(assert_type(i3.overlaps(pd.Interval(td1, td2, closed="left")), bool), bool) + + check(assert_type(i0 * 3, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1 * 3, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 * 3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 * 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i1 * 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 * 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(3 * i0, "pd.Interval[int]"), pd.Interval) + check(assert_type(3 * i1, "pd.Interval[float]"), pd.Interval) + check(assert_type(3 * i3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(3.5 * i0, "pd.Interval[float]"), pd.Interval) + check(assert_type(3.5 * i1, "pd.Interval[float]"), pd.Interval) + check(assert_type(3.5 * i3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 / 3, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1 / 3, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 / 3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 / 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i1 / 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 / 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 // 3, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1 // 3, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 // 3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 // 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i1 // 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 // 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 - 1, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1 - 1, "pd.Interval[float]"), pd.Interval) + check( + assert_type(i2 - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + ) + check( + assert_type(i3 - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + ) + + check(assert_type(i0 - 1.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i1 - 1.5, "pd.Interval[float]"), pd.Interval) + check( + assert_type(i2 - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + ) + check( + assert_type(i3 - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + ) + + check(assert_type(i0 + 1, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1 + 1, "pd.Interval[float]"), pd.Interval) + check( + assert_type(i2 + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + ) + check( + assert_type(i3 + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + ) + + check(assert_type(i0 + 1.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i1 + 1.5, "pd.Interval[float]"), pd.Interval) + check( + assert_type(i2 + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + ) + check( + assert_type(i3 + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + ) + + check(assert_type(0.5 in i0, bool), bool) + check(assert_type(1 in i0, bool), bool) + check(assert_type(1 in i1, bool), bool) + check(assert_type(pd.Timestamp("2000-1-1") in i2, bool), bool) + check(assert_type(pd.Timedelta(days=1) in i3, bool), bool) + + check(assert_type(hash(i0), int), int) + check(assert_type(hash(i1), int), int) + check(assert_type(hash(i2), int), int) + check(assert_type(hash(i3), int), int) + + +def test_timestamp() -> None: + + pd.Timestamp("2000-1-1") + pd.Timestamp("2000-1-1", tz="US/Pacific") + pd.Timestamp("2000-1-1", tz=pytz.timezone("US/Eastern")) + pd.Timestamp("2000-1-1", tz=dateutil.tz.UTC) + pd.Timestamp( + year=2000, + month=1, + day=1, + hour=1, + minute=1, + second=1, + microsecond=1, + nanosecond=1, + ) + pd.Timestamp(1, unit="D") + pd.Timestamp(1, unit="h") + pd.Timestamp(1, unit="m") + pd.Timestamp(1, unit="s") + pd.Timestamp(1, unit="ms") + pd.Timestamp(1, unit="us") + pd.Timestamp(year=2000, month=3, day=24, hour=12, minute=27, fold=0) + pd.Timestamp(year=2000, month=3, day=24, hour=12, minute=27, fold=1) + pd.Timestamp( + year=2000, + month=1, + day=1, + hour=1, + minute=1, + second=1, + microsecond=1, + nanosecond=1, + tzinfo=dt.timezone(offset=dt.timedelta(hours=6), name="EST"), + ) + ts = pd.Timestamp(year=2000, month=3, day=24, hour=12, minute=27) + check(assert_type(ts.asm8, np.datetime64), np.datetime64) + check(assert_type(ts.day_of_week, int), int) + check(assert_type(ts.day_of_year, int), int) + check(assert_type(ts.dayofweek, int), int) + check(assert_type(ts.dayofyear, int), int) + check(assert_type(ts.days_in_month, int), int) + check(assert_type(ts.daysinmonth, int), int) + check(assert_type(ts.is_leap_year, bool), bool) + check(assert_type(ts.is_month_end, bool), bool) + check(assert_type(ts.is_month_start, bool), bool) + check(assert_type(ts.is_quarter_end, bool), bool) + check(assert_type(ts.is_quarter_start, bool), bool) + check(assert_type(ts.is_year_end, bool), bool) + check(assert_type(ts.is_year_start, bool), bool) + check(assert_type(ts.quarter, int), int) + check(assert_type(ts.tz, Optional[dt.tzinfo]), type(None)) + check(assert_type(ts.week, int), int) + check(assert_type(ts.weekofyear, int), int) + check(assert_type(ts.day, int), int) + check(assert_type(ts.fold, int), int) + check(assert_type(ts.hour, int), int) + check(assert_type(ts.microsecond, int), int) + check(assert_type(ts.minute, int), int) + check(assert_type(ts.month, int), int) + check(assert_type(ts.nanosecond, int), int) + check(assert_type(ts.second, int), int) + check(assert_type(ts.tzinfo, Optional[dt.tzinfo]), type(None)) + check(assert_type(ts.value, int), int) + check(assert_type(ts.year, int), int) + + ts = pd.Timestamp("2000-1-1") + check(assert_type(ts + pd.Timedelta(days=1), pd.Timestamp), pd.Timestamp) + check(assert_type(ts + dt.timedelta(days=1), pd.Timestamp), pd.Timestamp) + check(assert_type(pd.Timedelta(days=1) + ts, pd.Timestamp), pd.Timestamp) + check(assert_type(dt.timedelta(days=1) + ts, pd.Timestamp), pd.Timestamp) + check(assert_type(ts + 3 * Day(), pd.Timestamp), pd.Timestamp) + check(assert_type(3 * Day() + ts, pd.Timestamp), pd.Timestamp) + check( + assert_type(ts + pd.TimedeltaIndex([1, 2, 3], "D"), pd.DatetimeIndex), + pd.DatetimeIndex, + ) + check( + assert_type(ts + pd.Series([1, 2], dtype="timedelta64[ns]"), "TimestampSeries"), + pd.Series, + ) + np_td64_arr: npt.NDArray[np.timedelta64] = np.array([1, 2], dtype="timedelta64[ns]") + np_dt64_arr: npt.NDArray[np.datetime64] = np.array( + [1, 2, 3], dtype="datetime64[ns]" + ) + check( + assert_type(ts + np_td64_arr, npt.NDArray[np.datetime64]), + np.ndarray, + ) + check( + assert_type(pd.TimedeltaIndex([1, 2, 3], "D") + ts, pd.DatetimeIndex), + pd.DatetimeIndex, + ) + check( + assert_type(pd.Series([1, 2], dtype="timedelta64[ns]") + ts, "TimestampSeries"), + pd.Series, + ) + + check(assert_type(ts - pd.Timedelta(days=1), pd.Timestamp), pd.Timestamp) + check(assert_type(ts - dt.timedelta(days=1), pd.Timestamp), pd.Timestamp) + check(assert_type(ts - 3 * Day(), pd.Timestamp), pd.Timestamp) + check( + assert_type(ts - pd.TimedeltaIndex([1, 2, 3], "D"), pd.DatetimeIndex), + pd.DatetimeIndex, + ) + ts_series = cast("TimedeltaSeries", pd.Series([1, 2], dtype="timedelta64[ns]")) + check( + assert_type(ts - ts_series, "TimestampSeries"), + pd.Series, + ) + check(assert_type(ts - np_td64_arr, npt.NDArray[np.datetime64]), np.ndarray) + + check(assert_type(ts > ts, bool), bool) + check(assert_type(ts > np.datetime64(1, "ns"), bool), bool) + check(assert_type(ts > dt.datetime(year=2000, month=1, day=1), bool), bool) + check(assert_type(ts > pd.DatetimeIndex(["2000-1-1"]), np_ndarray_bool), np.ndarray) + check(assert_type(ts > np_dt64_arr, np_ndarray_bool), np.ndarray) + check( + assert_type( + ts > pd.Series([1, 2, 3], dtype="datetime64[ns]"), "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(dt.datetime(year=2000, month=1, day=1) > ts, bool), bool) + check(assert_type(pd.DatetimeIndex(["2000-1-1"]) > ts, np_ndarray_bool), np.ndarray) + check(assert_type(np_dt64_arr > ts, np_ndarray_bool), np.ndarray) + check( + assert_type( + pd.Series([1, 2, 3], dtype="datetime64[ns]") > ts, "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(ts >= ts, bool), bool) + check(assert_type(ts >= np.datetime64(1, "ns"), bool), bool) + check(assert_type(ts >= dt.datetime(year=2000, month=1, day=1), bool), bool) + check( + assert_type(ts >= pd.DatetimeIndex(["2000-1-1"]), np_ndarray_bool), np.ndarray + ) + check(assert_type(ts >= np_dt64_arr, np_ndarray_bool), np.ndarray) + check( + assert_type( + ts >= pd.Series([1, 2, 3], dtype="datetime64[ns]"), "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(dt.datetime(year=2000, month=1, day=1) >= ts, bool), bool) + check( + assert_type(pd.DatetimeIndex(["2000-1-1"]) >= ts, np_ndarray_bool), np.ndarray + ) + check(assert_type(np_dt64_arr >= ts, np_ndarray_bool), np.ndarray) + check( + assert_type( + pd.Series([1, 2, 3], dtype="datetime64[ns]") >= ts, "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(ts < ts, bool), bool) + check(assert_type(ts < np.datetime64(1, "ns"), bool), bool) + check(assert_type(ts < dt.datetime(year=2000, month=1, day=1), bool), bool) + check(assert_type(ts < pd.DatetimeIndex(["2000-1-1"]), np_ndarray_bool), np.ndarray) + check(assert_type(ts < np_dt64_arr, np_ndarray_bool), np.ndarray) + check( + assert_type( + ts < pd.Series([1, 2, 3], dtype="datetime64[ns]"), "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(dt.datetime(year=2000, month=1, day=1) < ts, bool), bool) + check(assert_type(pd.DatetimeIndex(["2000-1-1"]) < ts, np_ndarray_bool), np.ndarray) + check(assert_type(np_dt64_arr < ts, np_ndarray_bool), np.ndarray) + check( + assert_type( + pd.Series([1, 2, 3], dtype="datetime64[ns]") < ts, "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(ts <= ts, bool), bool) + check(assert_type(ts <= np.datetime64(1, "ns"), bool), bool) + check(assert_type(ts <= dt.datetime(year=2000, month=1, day=1), bool), bool) + check( + assert_type(ts <= pd.DatetimeIndex(["2000-1-1"]), np_ndarray_bool), np.ndarray + ) + check(assert_type(ts <= np_dt64_arr, np_ndarray_bool), np.ndarray) + check( + assert_type( + ts <= pd.Series([1, 2, 3], dtype="datetime64[ns]"), "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(dt.datetime(year=2000, month=1, day=1) <= ts, bool), bool) + check( + assert_type(pd.DatetimeIndex(["2000-1-1"]) <= ts, np_ndarray_bool), np.ndarray + ) + check(assert_type(np_dt64_arr <= ts, np_ndarray_bool), np.ndarray) + check( + assert_type( + pd.Series([1, 2, 3], dtype="datetime64[ns]") <= ts, "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(ts == ts, bool), bool) + check(assert_type(ts == np.datetime64(1, "ns"), bool), bool) + check(assert_type(ts == dt.datetime(year=2000, month=1, day=1), bool), bool) + check( + assert_type(ts == pd.DatetimeIndex(["2000-1-1"]), np_ndarray_bool), np.ndarray + ) + check( + assert_type(ts == np_dt64_arr, np_ndarray_bool), + np.ndarray, + ) + check( + assert_type( + ts == pd.Series([1, 2, 3], dtype="datetime64[ns]"), "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(dt.datetime(year=2000, month=1, day=1) == ts, bool), bool) + check( + assert_type(pd.DatetimeIndex(["2000-1-1"]) == ts, np_ndarray_bool), np.ndarray + ) + + check( + assert_type( + pd.Series([1, 2, 3], dtype="datetime64[ns]") == ts, "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(ts != ts, bool), bool) + check(assert_type(ts != np.datetime64(1, "ns"), bool), bool) + check(assert_type(ts != dt.datetime(year=2000, month=1, day=1), bool), bool) + check( + assert_type(ts != pd.DatetimeIndex(["2000-1-1"]), np_ndarray_bool), np.ndarray + ) + check(assert_type(ts != np_dt64_arr, np_ndarray_bool), np.ndarray) + check( + assert_type( + ts != pd.Series([1, 2, 3], dtype="datetime64[ns]"), "pd.Series[bool]" + ), + pd.Series, + ) + + check(assert_type(dt.datetime(year=2000, month=1, day=1) != ts, bool), bool) + check( + assert_type(pd.DatetimeIndex(["2000-1-1"]) != ts, np_ndarray_bool), np.ndarray + ) + check( + assert_type( + pd.Series([1, 2, 3], dtype="datetime64[ns]") != ts, "pd.Series[bool]" + ), + pd.Series, + ) + + # Failures due to NumPy ops returning Any + check( + assert_type( # type: ignore[assert-type] + np_td64_arr + ts, npt.NDArray[np.datetime64] # type: ignore[operator] + ), + np.ndarray, + ) + check(assert_type(np.datetime64(1, "ns") > ts, np.bool_), np.bool_) # type: ignore[assert-type] + check(assert_type(np.datetime64(1, "ns") >= ts, np.bool_), np.bool_) # type: ignore[assert-type] + check(assert_type(np.datetime64(1, "ns") < ts, np.bool_), np.bool_) # type: ignore[assert-type] + check(assert_type(np.datetime64(1, "ns") <= ts, np.bool_), np.bool_) # type: ignore[assert-type] + check(assert_type(np.datetime64(1, "ns") == ts, np.bool_), np.bool_) # type: ignore[assert-type] + check(assert_type(np_dt64_arr == ts, np_ndarray_bool), np.ndarray) # type: ignore[assert-type] + check(assert_type(np.datetime64(1, "ns") != ts, np.bool_), np.bool_) # type: ignore[assert-type] + check(assert_type(np_dt64_arr != ts, np_ndarray_bool), np.ndarray) # type: ignore[assert-type] + + +def test_types_init() -> None: + ts: pd.Timestamp = pd.Timestamp("2021-03-01T12") + ts1: pd.Timestamp = pd.Timestamp(dt.date(2021, 3, 15)) + ts2: pd.Timestamp = pd.Timestamp(dt.datetime(2021, 3, 10, 12)) + ts3: pd.Timestamp = pd.Timestamp(pd.Timestamp("2021-03-01T12")) + ts4: pd.Timestamp = pd.Timestamp(1515590000.1, unit="s") + ts5: pd.Timestamp = pd.Timestamp(1515590000.1, unit="s", tz="US/Pacific") + ts6: pd.Timestamp = pd.Timestamp(1515590000100000000) # plain integer (nanosecond) + ts7: pd.Timestamp = pd.Timestamp(2021, 3, 10, 12) + ts8: pd.Timestamp = pd.Timestamp(year=2021, month=3, day=10, hour=12) + ts9: pd.Timestamp = pd.Timestamp( + year=2021, month=3, day=10, hour=12, tz="US/Pacific" + ) + + +def test_types_arithmetic() -> None: + ts: pd.Timestamp = pd.to_datetime("2021-03-01") + ts2: pd.Timestamp = pd.to_datetime("2021-01-01") + delta: pd.Timedelta = pd.to_timedelta("1 day") + + check(assert_type(ts - ts2, pd.Timedelta), pd.Timedelta) + check(assert_type(ts + delta, pd.Timestamp), pd.Timestamp) + check(assert_type(ts - delta, pd.Timestamp), pd.Timestamp) + check(assert_type(ts - dt.datetime(2021, 1, 3), pd.Timedelta), pd.Timedelta) + + +def test_types_comparison() -> None: + ts: pd.Timestamp = pd.to_datetime("2021-03-01") + ts2: pd.Timestamp = pd.to_datetime("2021-01-01") + + check(assert_type(ts < ts2, bool), bool) + check(assert_type(ts > ts2, bool), bool) + + +def test_types_timestamp_series_comparisons() -> None: + # GH 27 + df = pd.DataFrame(["2020-01-01", "2019-01-01"]) + tss = pd.to_datetime(df[0], format="%Y-%m-%d") + ts = pd.to_datetime("2019-02-01", format="%Y-%m-%d") + tssr = tss <= ts + tssr2 = tss >= ts + tssr3 = tss == ts + check(assert_type(tssr, "pd.Series[bool]"), pd.Series, bool) + check(assert_type(tssr2, "pd.Series[bool]"), pd.Series, bool) + check(assert_type(tssr3, "pd.Series[bool]"), pd.Series, bool) + # GH 265 + data = pd.date_range("2022-01-01", "2022-01-31", freq="D") + s = pd.Series(data) + ts2 = pd.Timestamp("2022-01-15") + check(assert_type(s, "TimestampSeries"), pd.Series, pd.Timestamp) + check(assert_type(ts2 <= s, "pd.Series[bool]"), pd.Series, bool) + check(assert_type(ts2 >= s, "pd.Series[bool]"), pd.Series, bool) + check(assert_type(ts2 < s, "pd.Series[bool]"), pd.Series, bool) + check(assert_type(ts2 > s, "pd.Series[bool]"), pd.Series, bool) + + +def test_types_pydatetime() -> None: + ts: pd.Timestamp = pd.Timestamp("2021-03-01T12") + check(assert_type(ts.to_pydatetime(), dt.datetime), dt.datetime) + check(assert_type(ts.to_pydatetime(False), dt.datetime), dt.datetime) + check(assert_type(ts.to_pydatetime(warn=True), dt.datetime), dt.datetime) + + +def test_timestamp_dateoffset_arithmetic() -> None: + ts = pd.Timestamp("2022-03-18") + do = pd.DateOffset(days=366) + check(assert_type(ts + do, pd.Timestamp), pd.Timestamp) + + +def test_todatetime_fromnumpy() -> None: + # GH 72 + t1 = np.datetime64("2022-07-04 02:30") + check(assert_type(pd.to_datetime(t1), pd.Timestamp), pd.Timestamp) diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index d43161ca0..8d15aac44 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -46,70 +46,6 @@ np_ndarray_bool = npt.NDArray[np.bool_] -def test_types_init() -> None: - ts: pd.Timestamp = pd.Timestamp("2021-03-01T12") - ts1: pd.Timestamp = pd.Timestamp(dt.date(2021, 3, 15)) - ts2: pd.Timestamp = pd.Timestamp(dt.datetime(2021, 3, 10, 12)) - ts3: pd.Timestamp = pd.Timestamp(pd.Timestamp("2021-03-01T12")) - ts4: pd.Timestamp = pd.Timestamp(1515590000.1, unit="s") - ts5: pd.Timestamp = pd.Timestamp(1515590000.1, unit="s", tz="US/Pacific") - ts6: pd.Timestamp = pd.Timestamp(1515590000100000000) # plain integer (nanosecond) - ts7: pd.Timestamp = pd.Timestamp(2021, 3, 10, 12) - ts8: pd.Timestamp = pd.Timestamp(year=2021, month=3, day=10, hour=12) - ts9: pd.Timestamp = pd.Timestamp( - year=2021, month=3, day=10, hour=12, tz="US/Pacific" - ) - - -def test_types_arithmetic() -> None: - ts: pd.Timestamp = pd.to_datetime("2021-03-01") - ts2: pd.Timestamp = pd.to_datetime("2021-01-01") - delta: pd.Timedelta = pd.to_timedelta("1 day") - - tsr: pd.Timedelta = ts - ts2 - tsr2: pd.Timestamp = ts + delta - tsr3: pd.Timestamp = ts - delta - tsr4: pd.Timedelta = ts - dt.datetime(2021, 1, 3) - - -def test_types_comparison() -> None: - ts: pd.Timestamp = pd.to_datetime("2021-03-01") - ts2: pd.Timestamp = pd.to_datetime("2021-01-01") - - tsr: bool = ts < ts2 - tsr2: bool = ts > ts2 - - -def test_types_timestamp_series_comparisons() -> None: - # GH 27 - df = pd.DataFrame(["2020-01-01", "2019-01-01"]) - tss = pd.to_datetime(df[0], format="%Y-%m-%d") - ts = pd.to_datetime("2019-02-01", format="%Y-%m-%d") - tssr = tss <= ts - tssr2 = tss >= ts - tssr3 = tss == ts - check(assert_type(tssr, "pd.Series[bool]"), pd.Series, bool) - check(assert_type(tssr2, "pd.Series[bool]"), pd.Series, bool) - check(assert_type(tssr3, "pd.Series[bool]"), pd.Series, bool) - # GH 265 - data = pd.date_range("2022-01-01", "2022-01-31", freq="D") - s = pd.Series(data) - ts2 = pd.Timestamp("2022-01-15") - check(assert_type(s, "TimestampSeries"), pd.Series, pd.Timestamp) - check(assert_type(ts2 <= s, "pd.Series[bool]"), pd.Series, bool) - check(assert_type(ts2 >= s, "pd.Series[bool]"), pd.Series, bool) - check(assert_type(ts2 < s, "pd.Series[bool]"), pd.Series, bool) - check(assert_type(ts2 > s, "pd.Series[bool]"), pd.Series, bool) - - -def test_types_pydatetime() -> None: - ts: pd.Timestamp = pd.Timestamp("2021-03-01T12") - - datet: dt.datetime = ts.to_pydatetime() - datet2: dt.datetime = ts.to_pydatetime(False) - datet3: dt.datetime = ts.to_pydatetime(warn=True) - - def test_to_timedelta() -> None: td: pd.Timedelta = pd.to_timedelta(3, "days") tds: pd.TimedeltaIndex = pd.to_timedelta([2, 3], "minutes") @@ -161,12 +97,6 @@ def test_timestamp_timedelta_series_arithmetic() -> None: check(assert_type(r8, "TimedeltaSeries"), pd.Series, pd.Timedelta) -def test_timestamp_dateoffset_arithmetic() -> None: - ts = pd.Timestamp("2022-03-18") - do = pd.DateOffset(days=366) - r1: pd.Timestamp = ts + do - - def test_datetimeindex_plus_timedelta() -> None: tscheck = pd.Series([pd.Timestamp("2022-03-05"), pd.Timestamp("2022-03-06")]) dti = pd.to_datetime(["2022-03-08", "2022-03-15"]) @@ -266,24 +196,6 @@ def test_dtindex_tzinfo() -> None: assert assert_type(dti.tzinfo, Optional[dt.tzinfo]) is None -def test_todatetime_fromnumpy() -> None: - # GH 72 - t1 = np.datetime64("2022-07-04 02:30") - check(assert_type(pd.to_datetime(t1), pd.Timestamp), pd.Timestamp) - - -def test_comparisons_datetimeindex() -> None: - # GH 74 - dti = pd.date_range("2000-01-01", "2000-01-10") - ts = pd.Timestamp("2000-01-05") - check(assert_type((dti < ts), np_ndarray_bool), np.ndarray) - check(assert_type((dti > ts), np_ndarray_bool), np.ndarray) - check(assert_type((dti >= ts), np_ndarray_bool), np.ndarray) - check(assert_type((dti <= ts), np_ndarray_bool), np.ndarray) - check(assert_type((dti == ts), np_ndarray_bool), np.ndarray) - check(assert_type((dti != ts), np_ndarray_bool), np.ndarray) - - def test_to_datetime_nat() -> None: # GH 88 check( @@ -663,6 +575,18 @@ def test_to_timedelta_scalar() -> None: ) +def test_comparisons_datetimeindex() -> None: + # GH 74 + dti = pd.date_range("2000-01-01", "2000-01-10") + ts = pd.Timestamp("2000-01-05") + check(assert_type((dti < ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti > ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti >= ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti <= ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti == ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti != ts), np_ndarray_bool), np.ndarray) + + def test_to_timedelta_series() -> None: s = pd.Series([10, 20, 30, 40]) s2 = pd.Series(["10ms", "20ms", "30ms"])