From a2cb27d715562484a8784bcc5b964257056a47a4 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Sun, 16 Oct 2022 13:19:42 +0100 Subject: [PATCH 01/10] ENH: Improve typing for Timedelta --- pandas-stubs/_libs/tslibs/period.pyi | 2 +- pandas-stubs/_libs/tslibs/timedeltas.pyi | 211 +++++++-- pandas-stubs/core/indexes/timedeltas.pyi | 11 + tests/test_scalars.py | 521 +++++++++++++++++++++++ 4 files changed, 713 insertions(+), 32 deletions(-) create mode 100644 tests/test_scalars.py diff --git a/pandas-stubs/_libs/tslibs/period.pyi b/pandas-stubs/_libs/tslibs/period.pyi index e0f93e54c..2225be6a0 100644 --- a/pandas-stubs/_libs/tslibs/period.pyi +++ b/pandas-stubs/_libs/tslibs/period.pyi @@ -6,7 +6,7 @@ class Period: def __init__( self, value: Any = ..., - freqstr: Any = ..., + freq: Any = ..., ordinal: Any = ..., year: Any = ..., month: int = ..., diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index aa09000ad..8cd8ebd65 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,179 @@ 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: Period) -> Period: ... + @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/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/tests/test_scalars.py b/tests/test_scalars.py new file mode 100644 index 000000000..6487a483f --- /dev/null +++ b/tests/test_scalars.py @@ -0,0 +1,521 @@ +from __future__ import annotations + +import datetime as dt +from typing import ( + TYPE_CHECKING, + Any, + cast, +) + +import numpy as np +from numpy import typing as npt +import pandas as pd +import pytest +from typing_extensions import assert_type + +from pandas._libs.tslibs.timedeltas import Components + +if TYPE_CHECKING: + 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_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_types_init() -> None: + check(assert_type(pd.Timestamp("2021-03-01T12"), pd.Timestamp), pd.Timestamp) + check(assert_type(pd.Timestamp(dt.date(2021, 3, 15)), pd.Timestamp), pd.Timestamp) + check( + assert_type(pd.Timestamp(dt.datetime(2021, 3, 10, 12)), pd.Timestamp), + pd.Timestamp, + ) + check( + assert_type(pd.Timestamp(pd.Timestamp("2021-03-01T12")), pd.Timestamp), + pd.Timestamp, + ) + check(assert_type(pd.Timestamp(1515590000.1, unit="s"), pd.Timestamp), pd.Timestamp) + check( + assert_type( + pd.Timestamp(1515590000.1, unit="s", tz="US/Pacific"), pd.Timestamp + ), + pd.Timestamp, + ) + check( + assert_type(pd.Timestamp(1515590000100000000), pd.Timestamp), pd.Timestamp + ) # plain integer (nanosecond) + check(assert_type(pd.Timestamp(2021, 3, 10, 12), pd.Timestamp), pd.Timestamp) + check( + assert_type(pd.Timestamp(year=2021, month=3, day=10, hour=12), pd.Timestamp), + pd.Timestamp, + ) + check( + assert_type( + pd.Timestamp(year=2021, month=3, day=10, hour=12, tz="US/Pacific"), + pd.Timestamp, + ), + pd.Timestamp, + ) + + +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) From e72654edfe74af55e47d11904b12afa645013541 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Mon, 17 Oct 2022 00:02:56 +0100 Subject: [PATCH 02/10] ENH: Improve typing of timedelta --- pandas-stubs/_libs/tslibs/timedeltas.pyi | 41 +++++++++++++++++++++--- pandas-stubs/core/series.pyi | 4 +-- tests/test_scalars.py | 40 +++++++++++++++++------ 3 files changed, 69 insertions(+), 16 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 8cd8ebd65..42f8490a7 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -10,6 +10,7 @@ from typing import ( ) import numpy as np +import pandas as pd from pandas import ( DatetimeIndex, Float64Index, @@ -132,6 +133,7 @@ class Timedelta(timedelta): def ceil(self: _S, freq: str | BaseOffset) -> _S: ... @property def resolution_string(self) -> str: ... + # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __add__(self, other: timedelta | np.timedelta64) -> Timedelta: ... @overload @@ -153,6 +155,10 @@ class Timedelta(timedelta): self, other: npt.NDArray[np.datetime64] ) -> npt.NDArray[np.datetime64]: ... @overload + def __add__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ... + @overload + def __add__(self, other: pd.Series) -> pd.Series: ... + @overload def __radd__(self, other: timedelta) -> Timedelta: ... @overload def __radd__( @@ -162,6 +168,9 @@ class Timedelta(timedelta): def __radd__( self, other: npt.NDArray[np.datetime64] ) -> npt.NDArray[np.datetime64]: ... + @overload + def __radd__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ... + # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __sub__(self, other: timedelta | np.timedelta64) -> Timedelta: ... @overload @@ -169,6 +178,10 @@ class Timedelta(timedelta): self, other: npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.timedelta64]: ... @overload + def __sub__(self, other: pd.TimedeltaIndex) -> TimedeltaIndex: ... + @overload + def __sub__(self, other: pd.Series) -> pd.Series: ... + @overload def __rsub__(self, other: timedelta | np.timedelta64) -> Timedelta: ... @overload def __rsub__(self, other: Timestamp | np.datetime64) -> Timestamp: ... @@ -186,15 +199,20 @@ class Timedelta(timedelta): def __rsub__( self, other: npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.timedelta64]: ... + @overload + def __rsub__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ... def __neg__(self) -> Timedelta: ... def __pos__(self) -> Timedelta: ... def __abs__(self) -> Timedelta: ... + # Override due to more types supported than dt.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: ... + def __mul__(self, other: Series[int]) -> TimedeltaSeries: ... + @overload + def __mul__(self, other: Series[float]) -> TimedeltaSeries: ... @overload def __mul__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... @overload @@ -202,9 +220,12 @@ class Timedelta(timedelta): @overload def __rmul__(self, other: np.ndarray) -> np.ndarray: ... @overload - def __rmul__(self, other: Series) -> Series: ... + def __rmul__(self, other: Series[int]) -> TimedeltaSeries: ... + @overload + def __rmul__(self, other: Series[float]) -> TimedeltaSeries: ... @overload def __rmul__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... + # Override due to more types supported than dt.timedelta # error: Signature of "__floordiv__" incompatible with supertype "timedelta" @overload # type: ignore[override] def __floordiv__(self, other: timedelta) -> int: ... @@ -221,13 +242,16 @@ class Timedelta(timedelta): @overload def __floordiv__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... @overload - def __floordiv__(self, other: Series) -> Series: ... + def __floordiv__(self, other: Series[int]) -> TimedeltaSeries: ... + @overload + def __floordiv__(self, other: Series[float]) -> TimedeltaSeries: ... @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]: ... + # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __truediv__(self, other: timedelta) -> float: ... @overload @@ -239,9 +263,12 @@ class Timedelta(timedelta): @overload def __truediv__(self, other: TimedeltaSeries) -> Series[float]: ... @overload - def __truediv__(self, other: Series) -> Series: ... + def __truediv__(self, other: Series[int]) -> TimedeltaSeries: ... + @overload + def __truediv__(self, other: Series[float]) -> TimedeltaSeries: ... @overload def __truediv__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... + # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __eq__(self, other: timedelta | np.timedelta64) -> bool: ... @overload @@ -250,6 +277,7 @@ class Timedelta(timedelta): def __eq__( self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... + # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __ne__(self, other: timedelta | np.timedelta64) -> bool: ... @overload @@ -258,6 +286,7 @@ class Timedelta(timedelta): def __ne__( self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... + # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __mod__(self, other: timedelta) -> Timedelta: ... @overload @@ -273,24 +302,28 @@ class Timedelta(timedelta): def __divmod__(self, other: timedelta) -> tuple[int, Timedelta]: ... # Mypy complains Forward operator "" is not callable, so ignore misc # for le, lt ge and gt + # Override due to more types supported than dt.timedelta @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_]: ... + # Override due to more types supported than dt.timedelta @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_]: ... + # Override due to more types supported than dt.timedelta @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_]: ... + # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __gt__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @overload diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 94161ed2d..9d07738d2 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1274,7 +1274,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): ) -> Series[int]: ... def __rsub__(self, other: num | _ListLike | Series[S1]) -> Series: ... @overload - def __rtruediv__(self, other: Timedelta | TimedeltaSeries) -> Series[float]: ... + def __rtruediv__(self, other: TimedeltaSeries) -> Series[float]: ... @overload def __rtruediv__(self, other: num | _ListLike | Series[S1]) -> Series: ... # ignore needed for mypy as we want different results based on the arguments @@ -1293,7 +1293,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def __sub__( self, other: Timedelta | TimedeltaSeries | TimedeltaIndex - ) -> TimestampSeries: ... + ) -> TimestampSeries | TimedeltaSeries: ... @overload def __sub__(self, other: num | _ListLike | Series) -> Series: ... @overload diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 6487a483f..359f5c6a4 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -4,6 +4,7 @@ from typing import ( TYPE_CHECKING, Any, + Union, cast, ) @@ -21,7 +22,7 @@ from pandas._typing import np_ndarray_bool else: - np_ndarray_bool = Any + TimedeltaSeries = TimestampSeries = np_ndarray_bool = Any from tests import check @@ -151,6 +152,10 @@ def test_timedelta() -> None: 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.TimedeltaIndex([td]), pd.TimedeltaIndex), pd.TimedeltaIndex + ) + check(assert_type(td + pd.Series([td]), pd.Series), pd.Series) check( assert_type( td + pd.period_range("2012-01-01", periods=3, freq="D"), pd.PeriodIndex @@ -175,6 +180,10 @@ def test_timedelta() -> None: ), np.ndarray, ) + check( + assert_type(pd.TimedeltaIndex([td]) + td, pd.TimedeltaIndex), pd.TimedeltaIndex + ) + check(assert_type(pd.Series([td]) + td, pd.Series), pd.Series) check(assert_type(td - td, pd.Timedelta), pd.Timedelta) check(assert_type(td - dt.timedelta(days=1), pd.Timedelta), pd.Timedelta) @@ -186,6 +195,10 @@ def test_timedelta() -> None: ), np.ndarray, ) + check( + assert_type(td - pd.TimedeltaIndex([td]), pd.TimedeltaIndex), pd.TimedeltaIndex + ) + check(assert_type(td - pd.Series([td]), pd.Series), pd.Series) # 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) @@ -211,6 +224,13 @@ def test_timedelta() -> None: assert_type(pd.date_range("2012-01-01", periods=3) - td, pd.DatetimeIndex), pd.DatetimeIndex, ) + check( + assert_type(pd.TimedeltaIndex([td]) - td, pd.TimedeltaIndex), pd.TimedeltaIndex + ) + check( + assert_type(pd.Series([td]) - td, Union[TimestampSeries, TimedeltaSeries]), + pd.Series, + ) check( assert_type( ndarray_td64 - td, # pyright: ignore[reportGeneralTypeIssues] @@ -234,15 +254,15 @@ def test_timedelta() -> None: 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 * pd.Series([1, 2, 3]), TimedeltaSeries), pd.Series) + check(assert_type(td * pd.Series([1.2, 2.2, 3.4]), TimedeltaSeries), 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(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) @@ -253,8 +273,8 @@ def test_timedelta() -> None: 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 // pd.Series([1, 2, 3]), TimedeltaSeries), pd.Series) + check(assert_type(td // pd.Series([1.2, 2.2, 3.4]), TimedeltaSeries), pd.Series) check(assert_type(td // i_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td // f_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) @@ -278,8 +298,8 @@ def test_timedelta() -> None: 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 / pd.Series([1, 2, 3]), TimedeltaSeries), pd.Series) + check(assert_type(td / pd.Series([1.2, 2.2, 3.4]), TimedeltaSeries), pd.Series) check(assert_type(td / i_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td / f_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) @@ -495,7 +515,7 @@ def test_types_timestamp_series_comparisons() -> None: 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(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) From 2bd346ed19801eed238e40e63403012286ad86dd Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Mon, 17 Oct 2022 10:57:51 +0100 Subject: [PATCH 03/10] ENH/TST: Improve Timedelta and its tests --- pandas-stubs/_libs/tslibs/timedeltas.pyi | 19 +- tests/test_scalars.py | 544 +++++++++++++---------- 2 files changed, 331 insertions(+), 232 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 42f8490a7..12fd79819 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -159,7 +159,9 @@ class Timedelta(timedelta): @overload def __add__(self, other: pd.Series) -> pd.Series: ... @overload - def __radd__(self, other: timedelta) -> Timedelta: ... + def __radd__(self, other: np.datetime64) -> Timestamp: ... + @overload + def __radd__(self, other: timedelta | np.timedelta64) -> Timedelta: ... @overload def __radd__( self, other: npt.NDArray[np.timedelta64] @@ -170,6 +172,8 @@ class Timedelta(timedelta): ) -> npt.NDArray[np.datetime64]: ... @overload def __radd__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ... + @overload + def __radd__(self, other: pd.PeriodIndex) -> pd.PeriodIndex: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __sub__(self, other: timedelta | np.timedelta64) -> Timedelta: ... @@ -272,7 +276,7 @@ class Timedelta(timedelta): @overload # type: ignore[override] def __eq__(self, other: timedelta | np.timedelta64) -> bool: ... @overload - def __eq__(self, other: Series) -> Series: ... + def __eq__(self, other: Series) -> Series[bool]: ... @overload def __eq__( self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] @@ -281,7 +285,7 @@ class Timedelta(timedelta): @overload # type: ignore[override] def __ne__(self, other: timedelta | np.timedelta64) -> bool: ... @overload - def __ne__(self, other: Series) -> Series: ... + def __ne__(self, other: Series) -> Series[bool]: ... @overload def __ne__( self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] @@ -309,6 +313,8 @@ class Timedelta(timedelta): def __le__( self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... + @overload + def __le__(self, other: Series) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __lt__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @@ -316,6 +322,8 @@ class Timedelta(timedelta): def __lt__( self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... + @overload + def __lt__(self, other: Series) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __ge__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @@ -323,6 +331,9 @@ class Timedelta(timedelta): def __ge__( self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... + @overload + def __ge__(self, other: Series) -> Series[bool]: ... + # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __gt__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @@ -330,6 +341,8 @@ class Timedelta(timedelta): def __gt__( self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... + @overload + def __gt__(self, other: Series) -> Series[bool]: ... def __hash__(self) -> int: ... def isoformat(self) -> str: ... def to_numpy(self) -> np.timedelta64: ... diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 359f5c6a4..008fb38fc 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -29,7 +29,7 @@ from pandas.tseries.offsets import Day -def test_timedelta() -> None: +def test_timedelta_construction() -> 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) @@ -115,6 +115,8 @@ def test_timedelta() -> None: check(assert_type(pd.Timedelta(weeks=1), pd.Timedelta), pd.Timedelta) check(assert_type(pd.Timedelta(milliseconds=1), pd.Timedelta), pd.Timedelta) + +def test_timedelta_properties_methods() -> None: td = pd.Timedelta("1 day") check(assert_type(td.value, int), int) check(assert_type(td.asm8, np.timedelta64), np.timedelta64) @@ -138,180 +140,201 @@ def test_timedelta() -> None: check(assert_type(td.view(np.int64), object), np.int64) check(assert_type(td.view("i8"), object), np.int64) + +def test_timedelta_add_sub() -> None: + td = pd.Timedelta("1 day") + 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.TimedeltaIndex([td]), pd.TimedeltaIndex), pd.TimedeltaIndex - ) - check(assert_type(td + pd.Series([td]), pd.Series), pd.Series) + as1 = pd.Period("2012-01-01", freq="D") + as2 = pd.Timestamp("2012-01-01") + as3 = dt.datetime(2012, 1, 1) + as4 = dt.date(2012, 1, 1) + as5 = np.datetime64(1, "ns") + as6 = dt.timedelta(days=1) + as7 = np.timedelta64(1, "D") + as8 = pd.TimedeltaIndex([td]) + as9 = pd.Series([td]) + as10 = pd.period_range("2012-01-01", periods=3, freq="D") + as11 = pd.date_range("2012-01-01", periods=3) + as12 = ndarray_td64 + as13 = ndarray_dt64 + + check(assert_type(td + td, pd.Timedelta), pd.Timedelta) + check(assert_type(td + as1, pd.Period), pd.Period) + check(assert_type(td + as2, pd.Timestamp), pd.Timestamp) + check(assert_type(td + as3, pd.Timestamp), pd.Timestamp) + check(assert_type(td + as4, dt.date), dt.date) + check(assert_type(td + as5, pd.Timestamp), pd.Timestamp) + check(assert_type(td + as6, pd.Timedelta), pd.Timedelta) + check(assert_type(td + as7, pd.Timedelta), pd.Timedelta) + check(assert_type(td + as8, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td + as9, pd.Series), pd.Series) + check(assert_type(td + as10, pd.PeriodIndex), pd.PeriodIndex) + check(assert_type(td + as11, pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(td + as12, npt.NDArray[np.timedelta64]), np.ndarray) + # pyright has trouble with timedelta64 and datetime64 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] + td + as13, # pyright: ignore[reportGeneralTypeIssues] npt.NDArray[np.datetime64], ), np.ndarray, ) + + check(assert_type(as1 + td, pd.Period), pd.Period) + check(assert_type(as2 + td, pd.Timestamp), pd.Timestamp) + check(assert_type(as3 + td, dt.datetime), dt.datetime) + check(assert_type(as4 + td, dt.date), dt.date) + check(assert_type(as5 + td, pd.Timestamp), pd.Timestamp) + # pyright is wrong here because as6 + td calls td.__radd__(as6), + # not timedelta.__add__ check( - assert_type(pd.TimedeltaIndex([td]) + td, pd.TimedeltaIndex), pd.TimedeltaIndex + assert_type(as6 + td, pd.Timedelta), # pyright: ignore[reportGeneralTypeIssues] + pd.Timedelta, ) - check(assert_type(pd.Series([td]) + td, pd.Series), pd.Series) - - 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(as7 + td, pd.Timedelta), pd.Timedelta) + check(assert_type(as8 + td, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(as9 + td, pd.Series), pd.Series) + check(assert_type(as10 + td, pd.PeriodIndex), pd.PeriodIndex) + check(assert_type(as11 + td, pd.DatetimeIndex), pd.DatetimeIndex) + # pyright is wrong here because ndarray.__add__(Timedelta) is NotImplemented check( assert_type( - td - ndarray_td64, + as12 + td, # pyright: ignore[reportGeneralTypeIssues] npt.NDArray[np.timedelta64], ), np.ndarray, ) - check( - assert_type(td - pd.TimedeltaIndex([td]), pd.TimedeltaIndex), pd.TimedeltaIndex - ) - check(assert_type(td - pd.Series([td]), pd.Series), pd.Series) - # 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) + + # sub is not symmetric with dates. In general date_like - timedelta is + # sensible, while timedelta - date_like is not + # TypeError: as1, as2, as3, as4, as5, as10, as11, as13 + check(assert_type(td - td, pd.Timedelta), pd.Timedelta) + check(assert_type(td - as6, pd.Timedelta), pd.Timedelta) + check(assert_type(td - as7, pd.Timedelta), pd.Timedelta) + check(assert_type(td - as8, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td - as9, pd.Series), pd.Series) + check(assert_type(td - as12, npt.NDArray[np.timedelta64]), np.ndarray) + + check(assert_type(as1 - td, pd.Period), pd.Period) + check(assert_type(as2 - td, pd.Timestamp), pd.Timestamp) + check(assert_type(as3 - td, dt.datetime), dt.datetime) + check(assert_type(as4 - td, dt.date), dt.date) + check(assert_type(as5 - td, pd.Timestamp), pd.Timestamp) + # pyright is wrong here because as6 + td calls td.__rsub__(as6), + # not timedelta.__sub__ check( assert_type( - dt.timedelta(days=1) - td, # pyright: ignore[reportGeneralTypeIssues] + as6 - 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(pd.TimedeltaIndex([td]) - td, pd.TimedeltaIndex), pd.TimedeltaIndex - ) - check( - assert_type(pd.Series([td]) - td, Union[TimestampSeries, TimedeltaSeries]), - pd.Series, - ) + check(assert_type(as7 - td, pd.Timedelta), pd.Timedelta) + check(assert_type(as8 - td, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(as9 - td, Union[TimestampSeries, TimedeltaSeries]), pd.Series) + check(assert_type(as10 - td, pd.PeriodIndex), pd.PeriodIndex) + check(assert_type(as11 - td, pd.DatetimeIndex), pd.DatetimeIndex) + # pyright is wrong here because ndarray.__sub__(Timedelta) is NotImplemented check( assert_type( - ndarray_td64 - td, # pyright: ignore[reportGeneralTypeIssues] + as12 - td, # pyright: ignore[reportGeneralTypeIssues] npt.NDArray[np.timedelta64], ), np.ndarray, ) + # pyright is wrong here because ndarray.__sub__(Timedelta) is NotImplemented check( assert_type( - ndarray_dt64 - td, # pyright: ignore[reportGeneralTypeIssues] + as13 - td, # pyright: ignore[reportGeneralTypeIssues] npt.NDArray[np.datetime64], ), np.ndarray, ) + +def test_timedelta_mul_div() -> None: + td = pd.Timedelta("1 day") + 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]), TimedeltaSeries), pd.Series) - check(assert_type(td * pd.Series([1.2, 2.2, 3.4]), TimedeltaSeries), 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]) + np_float_arr: npt.NDArray[np.floating] = np.array([1.2, 2.2, 3.4]) + + md1 = 3 + md2 = 3.5 + md3 = np_intp_arr + md4 = np_float_arr + md5 = pd.Series([1, 2, 3]) + md6 = pd.Series([1.2, 2.2, 3.4]) + md7 = i_idx + md8 = f_idx + + check(assert_type(td * md1, pd.Timedelta), pd.Timedelta) + check(assert_type(td * md2, pd.Timedelta), pd.Timedelta) + check(assert_type(td * md3, np.ndarray), np.ndarray) + check(assert_type(td * md4, np.ndarray), np.ndarray) + check(assert_type(td * md5, TimedeltaSeries), pd.Series) + check(assert_type(td * md6, TimedeltaSeries), pd.Series) + check(assert_type(td * md7, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td * md8, pd.TimedeltaIndex), pd.TimedeltaIndex) + + check(assert_type(md1 * td, pd.Timedelta), pd.Timedelta) + check(assert_type(md2 * td, pd.Timedelta), pd.Timedelta) + # pyright is wrong here ndarray.__mul__(Timedelta0 is NotImplemented + check( + assert_type(md3 * td, np.ndarray), # pyright: ignore[reportGeneralTypeIssues] + np.ndarray, + ) + check( + assert_type(md4 * td, np.ndarray), # pyright: ignore[reportGeneralTypeIssues] + np.ndarray, + ) + check(assert_type(md5 * td, TimedeltaSeries), pd.Series) + check(assert_type(md6 * td, TimedeltaSeries), pd.Series) + check(assert_type(md7 * td, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(md8 * td, pd.TimedeltaIndex), pd.TimedeltaIndex) + 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]), TimedeltaSeries), pd.Series) - check(assert_type(td // pd.Series([1.2, 2.2, 3.4]), TimedeltaSeries), 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 // md1, pd.Timedelta), pd.Timedelta) + check(assert_type(td // md2, pd.Timedelta), pd.Timedelta) + check(assert_type(td // md3, npt.NDArray[np.timedelta64]), np.ndarray) + check(assert_type(td // md4, npt.NDArray[np.timedelta64]), np.ndarray) + check(assert_type(td // md5, TimedeltaSeries), pd.Series) + check(assert_type(td // md6, TimedeltaSeries), pd.Series) + check(assert_type(td // md7, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td // md8, 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) + # TypeError: md1, md2, md3, md4, md5, md6, md7, md8 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]), TimedeltaSeries), pd.Series) - check(assert_type(td / pd.Series([1.2, 2.2, 3.4]), TimedeltaSeries), 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 / md1, pd.Timedelta), pd.Timedelta) + check(assert_type(td / md2, pd.Timedelta), pd.Timedelta) + check(assert_type(td / md3, npt.NDArray[np.timedelta64]), np.ndarray) + check(assert_type(td / md4, npt.NDArray[np.timedelta64]), np.ndarray) + check(assert_type(td / md5, TimedeltaSeries), pd.Series) + check(assert_type(td / md6, TimedeltaSeries), pd.Series) + check(assert_type(td / md7, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td / md8, 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) + # TypeError: md1, md2, md3, md4, md5, md6, md7, md8 + + +def test_timedelta_mod_abs_unary() -> None: + td = pd.Timedelta("1 day") + + 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) @@ -332,121 +355,184 @@ def test_timedelta() -> None: pd.TimedeltaIndex, ) + # mypy and pyright reports dt.timedelta, even though __abs__ returns Timedelta + check(assert_type(abs(td), dt.timedelta), pd.Timedelta) 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, +def test_timedelta_cmp() -> None: + td = pd.Timedelta("1 day") + ndarray_td64: npt.NDArray[np.timedelta64] = np.array( + [1, 2, 3], dtype="timedelta64[D]" ) + c1 = td + c2 = dt.timedelta(days=1) + c3 = np.timedelta64(1, "D") + c4 = ndarray_td64 + c5 = pd.TimedeltaIndex([1, 2, 3], unit="D") + c6 = pd.Series([1, 2, 3], dtype="timedelta64[D]") + + check(assert_type(td < c1, bool), bool) + check(assert_type(td < c2, bool), bool) + check(assert_type(td < c3, bool), bool) + check(assert_type(td < c4, np_ndarray_bool), np.ndarray) + check(assert_type(c5 < td, np_ndarray_bool), np.ndarray) + check(assert_type(c2 < td, bool), bool) + check(assert_type(c4 < td, np_ndarray_bool), np.ndarray) + check(assert_type(c5 < td, np_ndarray_bool), np.ndarray) + + gt = check(assert_type(td > c1, bool), bool) + le = check(assert_type(td <= c1, bool), bool) + assert gt != le + + gt = check(assert_type(td > c2, bool), bool) + le = check(assert_type(td <= c2, bool), bool) + assert gt != le + + gt = check(assert_type(td > c3, bool), bool) + le = check(assert_type(td <= c3, bool), bool) + assert gt != le + + gt_a = check(assert_type(td > c4, np_ndarray_bool), np.ndarray) + le_a = check(assert_type(td <= c4, np_ndarray_bool), np.ndarray) + assert (gt_a != le_a).all() + + gt_a = check(assert_type(td > c5, np_ndarray_bool), np.ndarray) + le_a = check(assert_type(td <= c5, np_ndarray_bool), np.ndarray) + assert (gt_a != le_a).all() + + gt_s = check(assert_type(td > c6, "pd.Series[bool]"), pd.Series, bool) + le_s = check(assert_type(td <= c6, "pd.Series[bool]"), pd.Series, bool) + assert (gt_s != le_s).all() + + gt = check(assert_type(c2 > td, bool), bool) + le = check(assert_type(c2 <= td, bool), bool) + assert gt != le + + gt_b = check(assert_type(c3 > td, Any), np.bool_) + le_b = check(assert_type(c3 <= td, Any), np.bool_) + assert gt_b != le_b + + gt_a = check(assert_type(c4 > td, np_ndarray_bool), np.ndarray) + le_a = check(assert_type(c4 <= td, np_ndarray_bool), np.ndarray) + assert (gt_a != le_a).all() + + gt_a = check(assert_type(c5 > td, np_ndarray_bool), np.ndarray) + le_a = check(assert_type(c5 <= td, np_ndarray_bool), np.ndarray) + assert (gt_a != le_a).all() + + eq_s = check(assert_type(c6 > td, "pd.Series[bool]"), pd.Series, bool) + ne_s = check(assert_type(c6 <= td, "pd.Series[bool]"), pd.Series, bool) + assert (eq_s != ne_s).all() + + lt = check(assert_type(td < c1, bool), bool) + ge = check(assert_type(td >= c1, bool), bool) + assert lt != ge - 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, - ) + lt = check(assert_type(td < c2, bool), bool) + ge = check(assert_type(td >= c2, bool), bool) + assert lt != ge - 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, - ) + lt = check(assert_type(td < c3, bool), bool) + ge = check(assert_type(td >= c3, bool), bool) + assert lt != ge - 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) + lt_a = check(assert_type(td < c4, np_ndarray_bool), np.ndarray) + ge_a = check(assert_type(td >= c4, np_ndarray_bool), np.ndarray) + assert (lt_a != ge_a).all() - check(assert_type(td == 1, bool), bool) - check(assert_type(td == (3 + 2j), bool), bool) + lt_a = check(assert_type(td < c5, np_ndarray_bool), np.ndarray) + ge_a = check(assert_type(td >= c5, np_ndarray_bool), np.ndarray) + assert (lt_a != ge_a).all() - 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, + eq_s = check(assert_type(td < c6, "pd.Series[bool]"), pd.Series, bool) + ne_s = check(assert_type(td >= c6, "pd.Series[bool]"), pd.Series, bool) + assert (eq_s != ne_s).all() + + lt = check(assert_type(c2 < td, bool), bool) + ge = check(assert_type(c2 >= td, bool), bool) + assert lt != ge + + lt_b = check(assert_type(c3 < td, Any), np.bool_) + ge_b = check(assert_type(c3 >= td, Any), np.bool_) + assert lt_b != ge_b + + lt_a = check(assert_type(c4 < td, np_ndarray_bool), np.ndarray) + ge_a = check(assert_type(c4 >= td, np_ndarray_bool), np.ndarray) + assert (lt_a != ge_a).all() + + lt_a = check(assert_type(c5 < td, np_ndarray_bool), np.ndarray) + ge_a = check(assert_type(c5 >= td, np_ndarray_bool), np.ndarray) + assert (lt_a != ge_a).all() + + eq_s = check(assert_type(c6 < td, "pd.Series[bool]"), pd.Series, bool) + ne_s = check(assert_type(c6 >= td, "pd.Series[bool]"), pd.Series, bool) + assert (eq_s != ne_s).all() + + eq = check(assert_type(td == td, bool), bool) + ne = check(assert_type(td != td, bool), bool) + assert eq != ne + + eq = check(assert_type(td == c2, bool), bool) + ne = check(assert_type(td != c2, bool), bool) + assert eq != ne + + eq = check(assert_type(td == c3, bool), bool) + ne = check(assert_type(td != c3, bool), bool) + assert eq != ne + + eq_a = check(assert_type(td == c4, np_ndarray_bool), np.ndarray) + ne_a = check(assert_type(td != c4, np_ndarray_bool), np.ndarray) + assert (eq_a != ne_a).all() + + eq_a = check(assert_type(td == c5, np_ndarray_bool), np.ndarray) + ne_a = check(assert_type(td != c5, np_ndarray_bool), np.ndarray) + assert (eq_a != ne_a).all() + + eq_s = check(assert_type(td == c6, "pd.Series[bool]"), pd.Series, bool) + ne_s = check(assert_type(td != c6, "pd.Series[bool]"), pd.Series, bool) + assert (eq_s != ne_s).all() + + eq = check(assert_type(c2 == td, bool), bool) + ne = check(assert_type(c2 != td, bool), bool) + assert eq != ne + + eq = check(assert_type(c3 == td, Any), np.bool_) + ne = check(assert_type(c3 != td, Any), np.bool_) + assert eq != ne + + eq_a = check(assert_type(c4 == td, Any), np.ndarray) + ne_a = check(assert_type(c4 != td, Any), np.ndarray) + assert (eq_a != ne_a).all() + + eq_a = check(assert_type(c5 == td, np_ndarray_bool), np.ndarray) + ne_a = check(assert_type(c5 != td, np_ndarray_bool), np.ndarray) + assert (eq_a != ne_a).all() + + eq_s = check(assert_type(c6 == td, "pd.Series[bool]"), pd.Series, bool) + ne_s = check(assert_type(c6 != td, "pd.Series[bool]"), pd.Series, bool) + assert (eq_s != ne_s).all() + + eq = check(assert_type(td == 1, bool), bool) + ne = check(assert_type(td != 1, bool), bool) + assert eq != ne + + eq = check(assert_type(td == (3 + 2j), bool), bool) + ne = check(assert_type(td != (3 + 2j), bool), bool) + assert eq != ne + + eq_s = check( + assert_type(td == pd.Series([1, 2, 3]), "pd.Series[bool]"), pd.Series, bool ) - 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, + ne_s = check( + assert_type(td != pd.Series([1, 2, 3]), "pd.Series[bool]"), pd.Series, bool ) - 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_types_init() -> None: + assert (eq_s != ne_s).all() + + +def test_timestamp_types_init() -> None: check(assert_type(pd.Timestamp("2021-03-01T12"), pd.Timestamp), pd.Timestamp) check(assert_type(pd.Timestamp(dt.date(2021, 3, 15)), pd.Timestamp), pd.Timestamp) check( @@ -481,7 +567,7 @@ def test_types_init() -> None: ) -def test_types_arithmetic() -> None: +def test_timestamp_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") @@ -492,7 +578,7 @@ def test_types_arithmetic() -> None: check(assert_type(ts - dt.datetime(2021, 1, 3), pd.Timedelta), pd.Timedelta) -def test_types_comparison() -> None: +def test_timestamp_types_comparison() -> None: ts: pd.Timestamp = pd.to_datetime("2021-03-01") ts2: pd.Timestamp = pd.to_datetime("2021-01-01") @@ -522,7 +608,7 @@ def test_types_timestamp_series_comparisons() -> None: check(assert_type(ts2 > s, "pd.Series[bool]"), pd.Series, bool) -def test_types_pydatetime() -> None: +def test_timestamp_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) From de4faec92761aab581f9ab1adc5d75967c62e10d Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Mon, 17 Oct 2022 19:15:08 +0100 Subject: [PATCH 04/10] ENH: Final changes for timedelta --- pandas-stubs/_libs/tslibs/timedeltas.pyi | 40 +++++++++++++------- tests/test_scalars.py | 48 +++++++++++++++++++++--- 2 files changed, 69 insertions(+), 19 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 12fd79819..cd3ad2be8 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -139,6 +139,8 @@ class Timedelta(timedelta): @overload def __add__(self, other: dt.datetime | np.datetime64 | Timestamp) -> Timestamp: ... @overload + def __add__(self, other: NaTType) -> NaTType: ... + @overload def __add__(self, other: Period) -> Period: ... @overload def __add__(self, other: dt.date) -> dt.date: ... @@ -163,6 +165,8 @@ class Timedelta(timedelta): @overload def __radd__(self, other: timedelta | np.timedelta64) -> Timedelta: ... @overload + def __radd__(self, other: NaTType) -> NaTType: ... + @overload def __radd__( self, other: npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.timedelta64]: ... @@ -178,6 +182,8 @@ class Timedelta(timedelta): @overload # type: ignore[override] def __sub__(self, other: timedelta | np.timedelta64) -> Timedelta: ... @overload + def __sub__(self, other: NaTType) -> NaTType: ... + @overload def __sub__( self, other: npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.timedelta64]: ... @@ -190,6 +196,8 @@ class Timedelta(timedelta): @overload def __rsub__(self, other: Timestamp | np.datetime64) -> Timestamp: ... @overload + def __rsub__(self, other: NaTType) -> NaTType: ... + @overload def __rsub__(self, other: Period) -> Period: ... @overload def __rsub__(self, other: PeriodIndex) -> PeriodIndex: ... @@ -242,7 +250,7 @@ class Timedelta(timedelta): @overload def __floordiv__( self, other: npt.NDArray[np.timedelta64] - ) -> npt.NDArray[np.intp]: ... + ) -> npt.NDArray[np.number]: ... @overload def __floordiv__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... @overload @@ -250,14 +258,18 @@ class Timedelta(timedelta): @overload def __floordiv__(self, other: Series[float]) -> TimedeltaSeries: ... @overload + def __floordiv__(self, other: NaTType | None) -> float: ... + @overload def __rfloordiv__(self, other: timedelta | str) -> int: ... @overload - def __rfloordiv__(self, other: NaTType | None) -> NaTType: ... + def __rfloordiv__(self, other: NaTType | None) -> float: ... @overload - def __rfloordiv__(self, other: np.ndarray) -> npt.NDArray[np.timedelta64]: ... + def __rfloordiv__( + self, other: npt.NDArray[np.timedelta64] + ) -> npt.NDArray[np.number]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] - def __truediv__(self, other: timedelta) -> float: ... + def __truediv__(self, other: timedelta | NaTType) -> float: ... @overload def __truediv__(self, other: float) -> Timedelta: ... @overload @@ -272,11 +284,12 @@ class Timedelta(timedelta): def __truediv__(self, other: Series[float]) -> TimedeltaSeries: ... @overload def __truediv__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... + def __rtruediv__(self, other: timedelta | NaTType) -> float: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __eq__(self, other: timedelta | np.timedelta64) -> bool: ... @overload - def __eq__(self, other: Series) -> Series[bool]: ... + def __eq__(self, other: TimedeltaSeries) -> Series[bool]: ... @overload def __eq__( self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] @@ -285,7 +298,7 @@ class Timedelta(timedelta): @overload # type: ignore[override] def __ne__(self, other: timedelta | np.timedelta64) -> bool: ... @overload - def __ne__(self, other: Series) -> Series[bool]: ... + def __ne__(self, other: TimedeltaSeries) -> Series[bool]: ... @overload def __ne__( self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] @@ -296,13 +309,15 @@ class Timedelta(timedelta): @overload def __mod__(self, other: float) -> Timedelta: ... @overload + def __mod__(self, other: Series[int] | Series[float]) -> TimedeltaSeries: ... + @overload + def __mod__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... + @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]: ... # Mypy complains Forward operator "" is not callable, so ignore misc # for le, lt ge and gt @@ -314,7 +329,7 @@ class Timedelta(timedelta): self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... @overload - def __le__(self, other: Series) -> Series[bool]: ... + def __le__(self, other: TimedeltaSeries) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __lt__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @@ -323,7 +338,7 @@ class Timedelta(timedelta): self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... @overload - def __lt__(self, other: Series) -> Series[bool]: ... + def __lt__(self, other: TimedeltaSeries) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __ge__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @@ -332,8 +347,7 @@ class Timedelta(timedelta): self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... @overload - def __ge__(self, other: Series) -> Series[bool]: ... - + def __ge__(self, other: TimedeltaSeries) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __gt__(self, other: timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @@ -342,7 +356,7 @@ class Timedelta(timedelta): self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... @overload - def __gt__(self, other: Series) -> Series[bool]: ... + def __gt__(self, other: TimedeltaSeries) -> Series[bool]: ... def __hash__(self) -> int: ... def isoformat(self) -> str: ... def to_numpy(self) -> np.timedelta64: ... diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 008fb38fc..ea7ac2860 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -14,17 +14,23 @@ import pytest from typing_extensions import assert_type +from pandas._libs.tslibs import NaTType from pandas._libs.tslibs.timedeltas import Components if TYPE_CHECKING: - from pandas.core.series import TimedeltaSeries # noqa: F401 - from pandas.core.series import TimestampSeries # noqa: F401 + from pandas.core.series import ( + TimedeltaSeries, + TimestampSeries, + ) from pandas._typing import np_ndarray_bool else: TimedeltaSeries = TimestampSeries = np_ndarray_bool = Any -from tests import check +from tests import ( + TYPE_CHECKING_INVALID_USAGE, + check, +) from pandas.tseries.offsets import Day @@ -163,6 +169,7 @@ def test_timedelta_add_sub() -> None: as11 = pd.date_range("2012-01-01", periods=3) as12 = ndarray_td64 as13 = ndarray_dt64 + as14 = pd.NaT check(assert_type(td + td, pd.Timedelta), pd.Timedelta) check(assert_type(td + as1, pd.Period), pd.Period) @@ -185,6 +192,7 @@ def test_timedelta_add_sub() -> None: ), np.ndarray, ) + check(assert_type(td + as14, NaTType), NaTType) check(assert_type(as1 + td, pd.Period), pd.Period) check(assert_type(as2 + td, pd.Timestamp), pd.Timestamp) @@ -210,6 +218,7 @@ def test_timedelta_add_sub() -> None: ), np.ndarray, ) + check(assert_type(as14 + td, NaTType), NaTType) # sub is not symmetric with dates. In general date_like - timedelta is # sensible, while timedelta - date_like is not @@ -220,7 +229,7 @@ def test_timedelta_add_sub() -> None: check(assert_type(td - as8, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td - as9, pd.Series), pd.Series) check(assert_type(td - as12, npt.NDArray[np.timedelta64]), np.ndarray) - + check(assert_type(td - as14, NaTType), NaTType) check(assert_type(as1 - td, pd.Period), pd.Period) check(assert_type(as2 - td, pd.Timestamp), pd.Timestamp) check(assert_type(as3 - td, dt.datetime), dt.datetime) @@ -256,6 +265,7 @@ def test_timedelta_add_sub() -> None: ), np.ndarray, ) + check(assert_type(as14 - td, NaTType), NaTType) def test_timedelta_mul_div() -> None: @@ -303,6 +313,7 @@ def test_timedelta_mul_div() -> None: check(assert_type(md8 * td, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td // td, int), int) + check(assert_type(td // pd.NaT, float), float) check(assert_type(td // md1, pd.Timedelta), pd.Timedelta) check(assert_type(td // md2, pd.Timedelta), pd.Timedelta) check(assert_type(td // md3, npt.NDArray[np.timedelta64]), np.ndarray) @@ -312,10 +323,21 @@ def test_timedelta_mul_div() -> None: check(assert_type(td // md7, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td // md8, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(pd.NaT // td, float), float) # Note: None of the reverse floordiv work # TypeError: md1, md2, md3, md4, md5, md6, md7, md8 + if TYPE_CHECKING_INVALID_USAGE: + md1 // td # type: ignore[operator] + md2 // td # type: ignore[operator] + md3 // td # type: ignore[operator] + md4 // td # type: ignore[operator] + md5 // td # type: ignore[operator] + md6 // td # type: ignore[operator] + md7 // td # type: ignore[operator] + md8 // td # type: ignore[operator] check(assert_type(td / td, float), float) + check(assert_type(td / pd.NaT, float), float) check(assert_type(td / md1, pd.Timedelta), pd.Timedelta) check(assert_type(td / md2, pd.Timedelta), pd.Timedelta) check(assert_type(td / md3, npt.NDArray[np.timedelta64]), np.ndarray) @@ -325,8 +347,20 @@ def test_timedelta_mul_div() -> None: check(assert_type(td / md7, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td / md8, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(pd.NaT / td, float), float) # Note: None of the reverse truediv work # TypeError: md1, md2, md3, md4, md5, md6, md7, md8 + if TYPE_CHECKING_INVALID_USAGE: + md1 / td # type: ignore[operator] + md2 / td # type: ignore[operator] + md3 / td # type: ignore[operator] + md4 / td # type: ignore[operator] + # TODO: Series.__truediv__ says it supports Timedelta + # it does not, in general, except for TimedeltaSeries + # md5 / td # type: ignore[operator] + # md6 / td # type: ignore[operator] + md7 / td # type: ignore[operator] + md8 / td # type: ignore[operator] def test_timedelta_mod_abs_unary() -> None: @@ -346,8 +380,10 @@ def test_timedelta_mod_abs_unary() -> None: 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) + int_series = pd.Series([1, 2, 3], dtype=int) + float_series = pd.Series([1.2, 2.2, 3.4], dtype=float) + check(assert_type(td % int_series, TimedeltaSeries), pd.Series) + check(assert_type(td % float_series, TimedeltaSeries), pd.Series) check(assert_type(td % i_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) check( From 9841c636481221234cafb7c1df09635b9188b3f8 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Mon, 24 Oct 2022 19:17:57 +0100 Subject: [PATCH 05/10] ENH: Further improvements to timestamp --- pandas-stubs/_libs/tslibs/timedeltas.pyi | 1 - pandas-stubs/core/indexes/timedeltas.pyi | 20 +- tests/test_scalars.py | 498 ++++++++++++++--------- 3 files changed, 303 insertions(+), 216 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 3c9594529..960567c57 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -370,7 +370,6 @@ class Timedelta(timedelta): ) -> npt.NDArray[np.bool_]: ... @overload def __gt__(self, other: TimedeltaSeries) -> Series[bool]: ... - def __hash__(self) -> int: ... def isoformat(self) -> str: ... def to_numpy(self) -> np.timedelta64: ... @property diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index 9e5a87251..376d7fd86 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/pandas-stubs/core/indexes/timedeltas.pyi @@ -1,9 +1,12 @@ +import datetime as dt from typing import ( Hashable, Literal, + Sequence, overload, ) +import numpy as np from pandas import ( DateOffset, Period, @@ -20,6 +23,7 @@ from pandas._libs import ( ) from pandas._libs.tslibs import BaseOffset from pandas._typing import ( + AnyArrayLike, TimedeltaConvertibleTypes, num, ) @@ -27,24 +31,16 @@ from pandas._typing import ( class TimedeltaIndex(DatetimeTimedeltaMixin, TimedeltaIndexProperties): def __init__( self, - data=..., + data: AnyArrayLike + | list[str] + | Sequence[dt.timedelta | Timedelta | np.timedelta64 | float] = ..., unit: Literal["D", "h", "m", "s", "ms", "us", "ns"] = ..., freq: str | BaseOffset = ..., closed: object = ..., - dtype=..., + dtype: Literal[" None: @@ -131,6 +129,21 @@ def test_timedelta_construction() -> None: 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) + check( + assert_type( + pd.Timedelta( + days=1, + seconds=1, + microseconds=1, + minutes=1, + hours=1, + weeks=1, + milliseconds=1, + ), + pd.Timedelta, + ), + pd.Timedelta, + ) def test_timedelta_properties_methods() -> None: @@ -167,109 +180,116 @@ def test_timedelta_add_sub() -> None: ndarray_dt64: npt.NDArray[np.datetime64] = np.array( [1, 2, 3], dtype="datetime64[D]" ) - as1 = pd.Period("2012-01-01", freq="D") - as2 = pd.Timestamp("2012-01-01") - as3 = dt.datetime(2012, 1, 1) - as4 = dt.date(2012, 1, 1) - as5 = np.datetime64(1, "ns") - as6 = dt.timedelta(days=1) - as7 = np.timedelta64(1, "D") - as8 = pd.TimedeltaIndex([td]) - as9 = pd.Series(as8) - as10 = pd.period_range("2012-01-01", periods=3, freq="D") - as11 = pd.date_range("2012-01-01", periods=3) - as12 = ndarray_td64 - as13 = ndarray_dt64 - as14 = pd.NaT + as_period = pd.Period("2012-01-01", freq="D") + as_timestamp = pd.Timestamp("2012-01-01") + as_datetime = dt.datetime(2012, 1, 1) + as_date = dt.date(2012, 1, 1) + as_datetime64 = np.datetime64(1, "ns") + as_dt_timedelta = dt.timedelta(days=1) + as_timedelta64 = np.timedelta64(1, "D") + as_timedelta_index = pd.TimedeltaIndex([td]) + as_timedelta_series = pd.Series(as_timedelta_index) + as_period_index = pd.period_range("2012-01-01", periods=3, freq="D") + as_datetime_index = pd.date_range("2012-01-01", periods=3) + as_ndarray_td64 = ndarray_td64 + as_ndarray_dt64 = ndarray_dt64 + as_nat = pd.NaT check(assert_type(td + td, pd.Timedelta), pd.Timedelta) - check(assert_type(td + as1, pd.Period), pd.Period) - check(assert_type(td + as2, pd.Timestamp), pd.Timestamp) - check(assert_type(td + as3, pd.Timestamp), pd.Timestamp) - check(assert_type(td + as4, dt.date), dt.date) - check(assert_type(td + as5, pd.Timestamp), pd.Timestamp) - check(assert_type(td + as6, pd.Timedelta), pd.Timedelta) - check(assert_type(td + as7, pd.Timedelta), pd.Timedelta) - check(assert_type(td + as8, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(td + as9, TimedeltaSeries), pd.Series) - check(assert_type(td + as10, pd.PeriodIndex), pd.PeriodIndex) - check(assert_type(td + as11, pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(td + as_period, pd.Period), pd.Period) + check(assert_type(td + as_timestamp, pd.Timestamp), pd.Timestamp) + check(assert_type(td + as_datetime, pd.Timestamp), pd.Timestamp) + check(assert_type(td + as_date, dt.date), dt.date) + check(assert_type(td + as_datetime64, pd.Timestamp), pd.Timestamp) + check(assert_type(td + as_dt_timedelta, pd.Timedelta), pd.Timedelta) + check(assert_type(td + as_timedelta64, pd.Timedelta), pd.Timedelta) + check(assert_type(td + as_timedelta_index, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td + as_timedelta_series, TimedeltaSeries), pd.Series) + check(assert_type(td + as_period_index, pd.PeriodIndex), pd.PeriodIndex) + check(assert_type(td + as_datetime_index, pd.DatetimeIndex), pd.DatetimeIndex) check( - assert_type(td + as12, npt.NDArray[np.timedelta64]), np.ndarray, np.timedelta64 + assert_type(td + as_ndarray_td64, npt.NDArray[np.timedelta64]), + np.ndarray, + np.timedelta64, ) # pyright has trouble with timedelta64 and datetime64 check( assert_type( - td + as13, # pyright: ignore[reportGeneralTypeIssues] + td + as_ndarray_dt64, # pyright: ignore[reportGeneralTypeIssues] npt.NDArray[np.datetime64], ), np.ndarray, np.datetime64, ) - check(assert_type(td + as14, NaTType), NaTType) - - check(assert_type(as1 + td, pd.Period), pd.Period) - check(assert_type(as2 + td, pd.Timestamp), pd.Timestamp) - check(assert_type(as3 + td, dt.datetime), dt.datetime) - check(assert_type(as4 + td, dt.date), dt.date) - check(assert_type(as5 + td, pd.Timestamp), pd.Timestamp) - # pyright is wrong here because as6 + td calls td.__radd__(as6), + check(assert_type(td + as_nat, NaTType), NaTType) + + check(assert_type(as_period + td, pd.Period), pd.Period) + check(assert_type(as_timestamp + td, pd.Timestamp), pd.Timestamp) + check(assert_type(as_datetime + td, dt.datetime), dt.datetime) + check(assert_type(as_date + td, dt.date), dt.date) + check(assert_type(as_datetime64 + td, pd.Timestamp), pd.Timestamp) + # pyright is fails here because as6 + td calls td.__radd__(as6), # not timedelta.__add__ + # https://github.com/microsoft/pyright/issues/4088 check( - assert_type(as6 + td, pd.Timedelta), # pyright: ignore[reportGeneralTypeIssues] + assert_type( + as_dt_timedelta + td, pd.Timedelta + ), # pyright: ignore[reportGeneralTypeIssues] pd.Timedelta, ) - check(assert_type(as7 + td, pd.Timedelta), pd.Timedelta) - check(assert_type(as8 + td, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(as9 + td, TimedeltaSeries), pd.Series) - check(assert_type(as10 + td, pd.PeriodIndex), pd.PeriodIndex) - check(assert_type(as11 + td, pd.DatetimeIndex), pd.DatetimeIndex) - # pyright is wrong here because ndarray.__add__(Timedelta) is NotImplemented + check(assert_type(as_timedelta64 + td, pd.Timedelta), pd.Timedelta) + check(assert_type(as_timedelta_index + td, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(as_timedelta_series + td, TimedeltaSeries), pd.Series) + check(assert_type(as_period_index + td, pd.PeriodIndex), pd.PeriodIndex) + check(assert_type(as_datetime_index + td, pd.DatetimeIndex), pd.DatetimeIndex) + # pyright is fails here because ndarray.__add__(Timedelta) is NotImplemented check( assert_type( - as12 + td, # pyright: ignore[reportGeneralTypeIssues] + as_ndarray_td64 + td, # pyright: ignore[reportGeneralTypeIssues] npt.NDArray[np.timedelta64], ), np.ndarray, np.timedelta64, ) - check(assert_type(as14 + td, NaTType), NaTType) + check(assert_type(as_nat + td, NaTType), NaTType) # sub is not symmetric with dates. In general date_like - timedelta is # sensible, while timedelta - date_like is not # TypeError: as1, as2, as3, as4, as5, as10, as11, as13 check(assert_type(td - td, pd.Timedelta), pd.Timedelta) - check(assert_type(td - as6, pd.Timedelta), pd.Timedelta) - check(assert_type(td - as7, pd.Timedelta), pd.Timedelta) - check(assert_type(td - as8, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(td - as9, TimedeltaSeries), pd.Series) + check(assert_type(td - as_dt_timedelta, pd.Timedelta), pd.Timedelta) + check(assert_type(td - as_timedelta64, pd.Timedelta), pd.Timedelta) + check(assert_type(td - as_timedelta_index, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td - as_timedelta_series, TimedeltaSeries), pd.Series) check( - assert_type(td - as12, npt.NDArray[np.timedelta64]), np.ndarray, np.timedelta64 - ) - check(assert_type(td - as14, NaTType), NaTType) - check(assert_type(as1 - td, pd.Period), pd.Period) - check(assert_type(as2 - td, pd.Timestamp), pd.Timestamp) - check(assert_type(as3 - td, dt.datetime), dt.datetime) - check(assert_type(as4 - td, dt.date), dt.date) - check(assert_type(as5 - td, pd.Timestamp), pd.Timestamp) + assert_type(td - as_ndarray_td64, npt.NDArray[np.timedelta64]), + np.ndarray, + np.timedelta64, + ) + check(assert_type(td - as_nat, NaTType), NaTType) + check(assert_type(as_period - td, pd.Period), pd.Period) + check(assert_type(as_timestamp - td, pd.Timestamp), pd.Timestamp) + check(assert_type(as_datetime - td, dt.datetime), dt.datetime) + check(assert_type(as_date - td, dt.date), dt.date) + check(assert_type(as_datetime64 - td, pd.Timestamp), pd.Timestamp) # pyright is wrong here because as6 + td calls td.__rsub__(as6), # not timedelta.__sub__ check( assert_type( - as6 - td, # pyright: ignore[reportGeneralTypeIssues] + as_dt_timedelta - td, # pyright: ignore[reportGeneralTypeIssues] pd.Timedelta, ), pd.Timedelta, ) - check(assert_type(as7 - td, pd.Timedelta), pd.Timedelta) - check(assert_type(as8 - td, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(as9 - td, TimedeltaSeries), pd.Series) - check(assert_type(as10 - td, pd.PeriodIndex), pd.PeriodIndex) - check(assert_type(as11 - td, pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(as_timedelta64 - td, pd.Timedelta), pd.Timedelta) + check(assert_type(as_timedelta_index - td, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(as_timedelta_series - td, TimedeltaSeries), pd.Series) + check(assert_type(as_period_index - td, pd.PeriodIndex), pd.PeriodIndex) + check(assert_type(as_datetime_index - td, pd.DatetimeIndex), pd.DatetimeIndex) # pyright is wrong here because ndarray.__sub__(Timedelta) is NotImplemented check( assert_type( - as12 - td, # pyright: ignore[reportGeneralTypeIssues] + as_ndarray_td64 - td, # pyright: ignore[reportGeneralTypeIssues] npt.NDArray[np.timedelta64], ), np.ndarray, @@ -278,13 +298,13 @@ def test_timedelta_add_sub() -> None: # pyright is wrong here because ndarray.__sub__(Timedelta) is NotImplemented check( assert_type( - as13 - td, # pyright: ignore[reportGeneralTypeIssues] + as_ndarray_dt64 - td, # pyright: ignore[reportGeneralTypeIssues] npt.NDArray[np.datetime64], ), np.ndarray, np.datetime64, ) - check(assert_type(as14 - td, NaTType), NaTType) + check(assert_type(as_nat - td, NaTType), NaTType) def test_timedelta_mul_div() -> None: @@ -297,89 +317,93 @@ def test_timedelta_mul_div() -> None: np_intp_arr: npt.NDArray[np.integer] = np.array([1, 2, 3]) np_float_arr: npt.NDArray[np.floating] = np.array([1.2, 2.2, 3.4]) - md1 = 3 - md2 = 3.5 - md3 = np_intp_arr - md4 = np_float_arr - md5 = pd.Series([1, 2, 3]) - md6 = pd.Series([1.2, 2.2, 3.4]) - md7 = i_idx - md8 = f_idx - - check(assert_type(td * md1, pd.Timedelta), pd.Timedelta) - check(assert_type(td * md2, pd.Timedelta), pd.Timedelta) - check(assert_type(td * md3, np.ndarray), np.ndarray) - check(assert_type(td * md4, np.ndarray), np.ndarray) - check(assert_type(td * md5, TimedeltaSeries), pd.Series) - check(assert_type(td * md6, TimedeltaSeries), pd.Series) - check(assert_type(td * md7, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(td * md8, pd.TimedeltaIndex), pd.TimedeltaIndex) - - check(assert_type(md1 * td, pd.Timedelta), pd.Timedelta) - check(assert_type(md2 * td, pd.Timedelta), pd.Timedelta) + md_int = 3 + md_float = 3.5 + md_ndarray_intp = np_intp_arr + md_ndarray_float = np_float_arr + mp_series_int = pd.Series([1, 2, 3]) + md_series_float = pd.Series([1.2, 2.2, 3.4]) + md_int64_index = i_idx + md_float_index = f_idx + + check(assert_type(td * md_int, pd.Timedelta), pd.Timedelta) + check(assert_type(td * md_float, pd.Timedelta), pd.Timedelta) + check(assert_type(td * md_ndarray_intp, np.ndarray), np.ndarray) + check(assert_type(td * md_ndarray_float, np.ndarray), np.ndarray) + check(assert_type(td * mp_series_int, TimedeltaSeries), pd.Series) + check(assert_type(td * md_series_float, TimedeltaSeries), pd.Series) + check(assert_type(td * md_int64_index, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td * md_float_index, pd.TimedeltaIndex), pd.TimedeltaIndex) + + check(assert_type(md_int * td, pd.Timedelta), pd.Timedelta) + check(assert_type(md_float * td, pd.Timedelta), pd.Timedelta) # pyright is wrong here ndarray.__mul__(Timedelta0 is NotImplemented check( - assert_type(md3 * td, np.ndarray), # pyright: ignore[reportGeneralTypeIssues] + assert_type( + md_ndarray_intp * td, np.ndarray + ), # pyright: ignore[reportGeneralTypeIssues] np.ndarray, ) check( - assert_type(md4 * td, np.ndarray), # pyright: ignore[reportGeneralTypeIssues] + assert_type( + md_ndarray_float * td, np.ndarray + ), # pyright: ignore[reportGeneralTypeIssues] np.ndarray, ) - check(assert_type(md5 * td, TimedeltaSeries), pd.Series) - check(assert_type(md6 * td, TimedeltaSeries), pd.Series) - check(assert_type(md7 * td, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(md8 * td, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(mp_series_int * td, TimedeltaSeries), pd.Series) + check(assert_type(md_series_float * td, TimedeltaSeries), pd.Series) + check(assert_type(md_int64_index * td, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(md_float_index * td, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td // td, int), int) check(assert_type(td // pd.NaT, float), float) - check(assert_type(td // md1, pd.Timedelta), pd.Timedelta) - check(assert_type(td // md2, pd.Timedelta), pd.Timedelta) - check(assert_type(td // md3, npt.NDArray[np.timedelta64]), np.ndarray) - check(assert_type(td // md4, npt.NDArray[np.timedelta64]), np.ndarray) - check(assert_type(td // md5, TimedeltaSeries), pd.Series) - check(assert_type(td // md6, TimedeltaSeries), pd.Series) - check(assert_type(td // md7, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(td // md8, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td // md_int, pd.Timedelta), pd.Timedelta) + check(assert_type(td // md_float, pd.Timedelta), pd.Timedelta) + check(assert_type(td // md_ndarray_intp, npt.NDArray[np.timedelta64]), np.ndarray) + check(assert_type(td // md_ndarray_float, npt.NDArray[np.timedelta64]), np.ndarray) + check(assert_type(td // mp_series_int, TimedeltaSeries), pd.Series) + check(assert_type(td // md_series_float, TimedeltaSeries), pd.Series) + check(assert_type(td // md_int64_index, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td // md_float_index, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(pd.NaT // td, float), float) # Note: None of the reverse floordiv work # TypeError: md1, md2, md3, md4, md5, md6, md7, md8 if TYPE_CHECKING_INVALID_USAGE: - md1 // td # type: ignore[operator] - md2 // td # type: ignore[operator] - md3 // td # type: ignore[operator] - md4 // td # type: ignore[operator] - md5 // td # type: ignore[operator] - md6 // td # type: ignore[operator] - md7 // td # type: ignore[operator] - md8 // td # type: ignore[operator] + md_int // td # type: ignore[operator] + md_float // td # type: ignore[operator] + md_ndarray_intp // td # type: ignore[operator] + md_ndarray_float // td # type: ignore[operator] + mp_series_int // td # type: ignore[operator] + md_series_float // td # type: ignore[operator] + md_int64_index // td # type: ignore[operator] + md_float_index // td # type: ignore[operator] check(assert_type(td / td, float), float) check(assert_type(td / pd.NaT, float), float) - check(assert_type(td / md1, pd.Timedelta), pd.Timedelta) - check(assert_type(td / md2, pd.Timedelta), pd.Timedelta) - check(assert_type(td / md3, npt.NDArray[np.timedelta64]), np.ndarray) - check(assert_type(td / md4, npt.NDArray[np.timedelta64]), np.ndarray) - check(assert_type(td / md5, TimedeltaSeries), pd.Series) - check(assert_type(td / md6, TimedeltaSeries), pd.Series) - check(assert_type(td / md7, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(td / md8, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td / md_int, pd.Timedelta), pd.Timedelta) + check(assert_type(td / md_float, pd.Timedelta), pd.Timedelta) + check(assert_type(td / md_ndarray_intp, npt.NDArray[np.timedelta64]), np.ndarray) + check(assert_type(td / md_ndarray_float, npt.NDArray[np.timedelta64]), np.ndarray) + check(assert_type(td / mp_series_int, TimedeltaSeries), pd.Series) + check(assert_type(td / md_series_float, TimedeltaSeries), pd.Series) + check(assert_type(td / md_int64_index, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td / md_float_index, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(pd.NaT / td, float), float) # Note: None of the reverse truediv work # TypeError: md1, md2, md3, md4, md5, md6, md7, md8 if TYPE_CHECKING_INVALID_USAGE: - md1 / td # type: ignore[operator] - md2 / td # type: ignore[operator] - md3 / td # type: ignore[operator] - md4 / td # type: ignore[operator] + md_int / td # type: ignore[operator] + md_float / td # type: ignore[operator] + md_ndarray_intp / td # type: ignore[operator] + md_ndarray_float / td # type: ignore[operator] # TODO: Series.__truediv__ says it supports Timedelta # it does not, in general, except for TimedeltaSeries # md5 / td # type: ignore[operator] # md6 / td # type: ignore[operator] - md7 / td # type: ignore[operator] - md8 / td # type: ignore[operator] + md_int64_index / td # type: ignore[operator] + md_float_index / td # type: ignore[operator] def test_timedelta_mod_abs_unary() -> None: @@ -425,152 +449,220 @@ def test_timedelta_cmp() -> None: ndarray_td64: npt.NDArray[np.timedelta64] = np.array( [1, 2, 3], dtype="timedelta64[D]" ) - c1 = td - c2 = dt.timedelta(days=1) - c3 = np.timedelta64(1, "D") - c4 = ndarray_td64 - c5 = pd.TimedeltaIndex([1, 2, 3], unit="D") - c6 = pd.Series([1, 2, 3], dtype="timedelta64[D]") - - check(assert_type(td < c1, bool), bool) - check(assert_type(td < c2, bool), bool) - check(assert_type(td < c3, bool), bool) - check(assert_type(td < c4, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c5 < td, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c2 < td, bool), bool) - check(assert_type(c4 < td, np_ndarray_bool), np.ndarray, np.bool_) - check(assert_type(c5 < td, np_ndarray_bool), np.ndarray, np.bool_) - - gt = check(assert_type(td > c1, bool), bool) - le = check(assert_type(td <= c1, bool), bool) + c_timedelta = td + c_dt_timedelta = dt.timedelta(days=1) + c_timedelta64 = np.timedelta64(1, "D") + c_ndarray_td64 = ndarray_td64 + c_timedelta_index = pd.TimedeltaIndex([1, 2, 3], unit="D") + c_timedelta_series = pd.Series([1, 2, 3], dtype="timedelta64[D]") + + check(assert_type(td < c_timedelta, bool), bool) + check(assert_type(td < c_dt_timedelta, bool), bool) + check(assert_type(td < c_timedelta64, bool), bool) + check(assert_type(td < c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_) + check(assert_type(c_timedelta_index < td, np_ndarray_bool), np.ndarray, np.bool_) + check(assert_type(c_dt_timedelta < td, bool), bool) + check(assert_type(c_ndarray_td64 < td, np_ndarray_bool), np.ndarray, np.bool_) + check(assert_type(c_timedelta_index < td, np_ndarray_bool), np.ndarray, np.bool_) + + gt = check(assert_type(td > c_timedelta, bool), bool) + le = check(assert_type(td <= c_timedelta, bool), bool) assert gt != le - gt = check(assert_type(td > c2, bool), bool) - le = check(assert_type(td <= c2, bool), bool) + gt = check(assert_type(td > c_dt_timedelta, bool), bool) + le = check(assert_type(td <= c_dt_timedelta, bool), bool) assert gt != le - gt = check(assert_type(td > c3, bool), bool) - le = check(assert_type(td <= c3, bool), bool) + gt = check(assert_type(td > c_timedelta64, bool), bool) + le = check(assert_type(td <= c_timedelta64, bool), bool) assert gt != le - gt_a = check(assert_type(td > c4, np_ndarray_bool), np.ndarray, np.bool_) - le_a = check(assert_type(td <= c4, np_ndarray_bool), np.ndarray, np.bool_) + gt_a = check( + assert_type(td > c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ + ) + le_a = check( + assert_type(td <= c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (gt_a != le_a).all() - gt_a = check(assert_type(td > c5, np_ndarray_bool), np.ndarray, np.bool_) - le_a = check(assert_type(td <= c5, np_ndarray_bool), np.ndarray, np.bool_) + gt_a = check( + assert_type(td > c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ + ) + le_a = check( + assert_type(td <= c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (gt_a != le_a).all() - gt_s = check(assert_type(td > c6, "pd.Series[bool]"), pd.Series, bool) - le_s = check(assert_type(td <= c6, "pd.Series[bool]"), pd.Series, bool) + gt_s = check( + assert_type(td > c_timedelta_series, "pd.Series[bool]"), pd.Series, bool + ) + le_s = check( + assert_type(td <= c_timedelta_series, "pd.Series[bool]"), pd.Series, bool + ) assert (gt_s != le_s).all() - gt = check(assert_type(c2 > td, bool), bool) - le = check(assert_type(c2 <= td, bool), bool) + gt = check(assert_type(c_dt_timedelta > td, bool), bool) + le = check(assert_type(c_dt_timedelta <= td, bool), bool) assert gt != le - gt_b = check(assert_type(c3 > td, Any), np.bool_) - le_b = check(assert_type(c3 <= td, Any), np.bool_) + gt_b = check(assert_type(c_timedelta64 > td, Any), np.bool_) + le_b = check(assert_type(c_timedelta64 <= td, Any), np.bool_) assert gt_b != le_b - gt_a = check(assert_type(c4 > td, np_ndarray_bool), np.ndarray, np.bool_) - le_a = check(assert_type(c4 <= td, np_ndarray_bool), np.ndarray, np.bool_) + gt_a = check( + assert_type(c_ndarray_td64 > td, np_ndarray_bool), np.ndarray, np.bool_ + ) + le_a = check( + assert_type(c_ndarray_td64 <= td, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (gt_a != le_a).all() - gt_a = check(assert_type(c5 > td, np_ndarray_bool), np.ndarray, np.bool_) - le_a = check(assert_type(c5 <= td, np_ndarray_bool), np.ndarray, np.bool_) + gt_a = check( + assert_type(c_timedelta_index > td, np_ndarray_bool), np.ndarray, np.bool_ + ) + le_a = check( + assert_type(c_timedelta_index <= td, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (gt_a != le_a).all() - eq_s = check(assert_type(c6 > td, "pd.Series[bool]"), pd.Series, bool) - ne_s = check(assert_type(c6 <= td, "pd.Series[bool]"), pd.Series, bool) + eq_s = check( + assert_type(c_timedelta_series > td, "pd.Series[bool]"), pd.Series, bool + ) + ne_s = check( + assert_type(c_timedelta_series <= td, "pd.Series[bool]"), pd.Series, bool + ) assert (eq_s != ne_s).all() - lt = check(assert_type(td < c1, bool), bool) - ge = check(assert_type(td >= c1, bool), bool) + lt = check(assert_type(td < c_timedelta, bool), bool) + ge = check(assert_type(td >= c_timedelta, bool), bool) assert lt != ge - lt = check(assert_type(td < c2, bool), bool) - ge = check(assert_type(td >= c2, bool), bool) + lt = check(assert_type(td < c_dt_timedelta, bool), bool) + ge = check(assert_type(td >= c_dt_timedelta, bool), bool) assert lt != ge - lt = check(assert_type(td < c3, bool), bool) - ge = check(assert_type(td >= c3, bool), bool) + lt = check(assert_type(td < c_timedelta64, bool), bool) + ge = check(assert_type(td >= c_timedelta64, bool), bool) assert lt != ge - lt_a = check(assert_type(td < c4, np_ndarray_bool), np.ndarray, np.bool_) - ge_a = check(assert_type(td >= c4, np_ndarray_bool), np.ndarray, np.bool_) + lt_a = check( + assert_type(td < c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ + ) + ge_a = check( + assert_type(td >= c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (lt_a != ge_a).all() - lt_a = check(assert_type(td < c5, np_ndarray_bool), np.ndarray, np.bool_) - ge_a = check(assert_type(td >= c5, np_ndarray_bool), np.ndarray, np.bool_) + lt_a = check( + assert_type(td < c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ + ) + ge_a = check( + assert_type(td >= c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (lt_a != ge_a).all() - eq_s = check(assert_type(td < c6, "pd.Series[bool]"), pd.Series, bool) - ne_s = check(assert_type(td >= c6, "pd.Series[bool]"), pd.Series, bool) + eq_s = check( + assert_type(td < c_timedelta_series, "pd.Series[bool]"), pd.Series, bool + ) + ne_s = check( + assert_type(td >= c_timedelta_series, "pd.Series[bool]"), pd.Series, bool + ) assert (eq_s != ne_s).all() - lt = check(assert_type(c2 < td, bool), bool) - ge = check(assert_type(c2 >= td, bool), bool) + lt = check(assert_type(c_dt_timedelta < td, bool), bool) + ge = check(assert_type(c_dt_timedelta >= td, bool), bool) assert lt != ge - lt_b = check(assert_type(c3 < td, Any), np.bool_) - ge_b = check(assert_type(c3 >= td, Any), np.bool_) + lt_b = check(assert_type(c_timedelta64 < td, Any), np.bool_) + ge_b = check(assert_type(c_timedelta64 >= td, Any), np.bool_) assert lt_b != ge_b - lt_a = check(assert_type(c4 < td, np_ndarray_bool), np.ndarray, np.bool_) - ge_a = check(assert_type(c4 >= td, np_ndarray_bool), np.ndarray, np.bool_) + lt_a = check( + assert_type(c_ndarray_td64 < td, np_ndarray_bool), np.ndarray, np.bool_ + ) + ge_a = check( + assert_type(c_ndarray_td64 >= td, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (lt_a != ge_a).all() - lt_a = check(assert_type(c5 < td, np_ndarray_bool), np.ndarray, np.bool_) - ge_a = check(assert_type(c5 >= td, np_ndarray_bool), np.ndarray, np.bool_) + lt_a = check( + assert_type(c_timedelta_index < td, np_ndarray_bool), np.ndarray, np.bool_ + ) + ge_a = check( + assert_type(c_timedelta_index >= td, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (lt_a != ge_a).all() - eq_s = check(assert_type(c6 < td, "pd.Series[bool]"), pd.Series, bool) - ne_s = check(assert_type(c6 >= td, "pd.Series[bool]"), pd.Series, bool) + eq_s = check( + assert_type(c_timedelta_series < td, "pd.Series[bool]"), pd.Series, bool + ) + ne_s = check( + assert_type(c_timedelta_series >= td, "pd.Series[bool]"), pd.Series, bool + ) assert (eq_s != ne_s).all() eq = check(assert_type(td == td, bool), bool) ne = check(assert_type(td != td, bool), bool) assert eq != ne - eq = check(assert_type(td == c2, bool), bool) - ne = check(assert_type(td != c2, bool), bool) + eq = check(assert_type(td == c_dt_timedelta, bool), bool) + ne = check(assert_type(td != c_dt_timedelta, bool), bool) assert eq != ne - eq = check(assert_type(td == c3, bool), bool) - ne = check(assert_type(td != c3, bool), bool) + eq = check(assert_type(td == c_timedelta64, bool), bool) + ne = check(assert_type(td != c_timedelta64, bool), bool) assert eq != ne - eq_a = check(assert_type(td == c4, np_ndarray_bool), np.ndarray, np.bool_) - ne_a = check(assert_type(td != c4, np_ndarray_bool), np.ndarray, np.bool_) + eq_a = check( + assert_type(td == c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ + ) + ne_a = check( + assert_type(td != c_ndarray_td64, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (eq_a != ne_a).all() - eq_a = check(assert_type(td == c5, np_ndarray_bool), np.ndarray, np.bool_) - ne_a = check(assert_type(td != c5, np_ndarray_bool), np.ndarray, np.bool_) + eq_a = check( + assert_type(td == c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ + ) + ne_a = check( + assert_type(td != c_timedelta_index, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (eq_a != ne_a).all() - eq_s = check(assert_type(td == c6, "pd.Series[bool]"), pd.Series, bool) - ne_s = check(assert_type(td != c6, "pd.Series[bool]"), pd.Series, bool) + eq_s = check( + assert_type(td == c_timedelta_series, "pd.Series[bool]"), pd.Series, bool + ) + ne_s = check( + assert_type(td != c_timedelta_series, "pd.Series[bool]"), pd.Series, bool + ) assert (eq_s != ne_s).all() - eq = check(assert_type(c2 == td, bool), bool) - ne = check(assert_type(c2 != td, bool), bool) + eq = check(assert_type(c_dt_timedelta == td, bool), bool) + ne = check(assert_type(c_dt_timedelta != td, bool), bool) assert eq != ne - eq = check(assert_type(c3 == td, Any), np.bool_) - ne = check(assert_type(c3 != td, Any), np.bool_) + eq = check(assert_type(c_timedelta64 == td, Any), np.bool_) + ne = check(assert_type(c_timedelta64 != td, Any), np.bool_) assert eq != ne - eq_a = check(assert_type(c4 == td, Any), np.ndarray) - ne_a = check(assert_type(c4 != td, Any), np.ndarray) + eq_a = check(assert_type(c_ndarray_td64 == td, Any), np.ndarray) + ne_a = check(assert_type(c_ndarray_td64 != td, Any), np.ndarray) assert (eq_a != ne_a).all() - eq_a = check(assert_type(c5 == td, np_ndarray_bool), np.ndarray, np.bool_) - ne_a = check(assert_type(c5 != td, np_ndarray_bool), np.ndarray, np.bool_) + eq_a = check( + assert_type(c_timedelta_index == td, np_ndarray_bool), np.ndarray, np.bool_ + ) + ne_a = check( + assert_type(c_timedelta_index != td, np_ndarray_bool), np.ndarray, np.bool_ + ) assert (eq_a != ne_a).all() - eq_s = check(assert_type(c6 == td, "pd.Series[bool]"), pd.Series, bool) - ne_s = check(assert_type(c6 != td, "pd.Series[bool]"), pd.Series, bool) + eq_s = check( + assert_type(c_timedelta_series == td, "pd.Series[bool]"), pd.Series, bool + ) + ne_s = check( + assert_type(c_timedelta_series != td, "pd.Series[bool]"), pd.Series, bool + ) assert (eq_s != ne_s).all() eq = check(assert_type(td == 1, bool), bool) From 0aafaf24bfdc4edab61473f05f2a1771b93ecdb3 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Tue, 25 Oct 2022 01:07:20 +0100 Subject: [PATCH 06/10] ENH: Add more types --- pandas-stubs/_libs/tslibs/timedeltas.pyi | 40 +++++++++++++++--------- tests/test_scalars.py | 37 ++++++++++------------ 2 files changed, 41 insertions(+), 36 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 960567c57..767d7b4a9 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -163,7 +163,7 @@ class Timedelta(timedelta): def __add__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ... @overload def __add__( - self, other: Series[pd.Timedelta] | TimedeltaSeries + self, other: TimedeltaSeries | Series[pd.Timedelta] ) -> TimedeltaSeries: ... @overload def __add__( @@ -251,7 +251,7 @@ class Timedelta(timedelta): # Override due to more types supported than dt.timedelta # error: Signature of "__floordiv__" incompatible with supertype "timedelta" @overload # type: ignore[override] - def __floordiv__(self, other: timedelta) -> int: ... + def __floordiv__(self, other: timedelta | Timedelta | np.timedelta64) -> int: ... @overload def __floordiv__(self, other: float) -> Timedelta: ... @overload @@ -269,6 +269,10 @@ class Timedelta(timedelta): @overload def __floordiv__(self, other: Series[float]) -> TimedeltaSeries: ... @overload + def __floordiv__( + self, other: Series[Timedelta] | TimedeltaSeries + ) -> Series[int]: ... + @overload def __floordiv__(self, other: NaTType | None) -> float: ... @overload def __rfloordiv__(self, other: timedelta | Timedelta | str) -> int: ... @@ -288,7 +292,9 @@ class Timedelta(timedelta): self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating] ) -> npt.NDArray[np.timedelta64]: ... @overload - def __truediv__(self, other: TimedeltaSeries) -> Series[float]: ... + def __truediv__( + self, other: Series[Timedelta] | TimedeltaSeries + ) -> Series[float]: ... @overload def __truediv__(self, other: Series[int]) -> TimedeltaSeries: ... @overload @@ -297,23 +303,27 @@ class Timedelta(timedelta): def __truediv__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ... def __rtruediv__(self, other: timedelta | Timedelta | NaTType) -> float: ... # Override due to more types supported than dt.timedelta - @overload # type: ignore[override] - def __eq__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... @overload - def __eq__(self, other: TimedeltaSeries) -> Series[bool]: ... + def __eq__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @overload - def __eq__( + def __eq__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # type: ignore[misc] + @overload + def __eq__( # type: ignore[misc] self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... + @overload + def __eq__(self, other: object) -> Literal[False]: ... # Override due to more types supported than dt.timedelta - @overload # type: ignore[override] - def __ne__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... @overload - def __ne__(self, other: TimedeltaSeries) -> Series[bool]: ... + def __ne__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @overload - def __ne__( + def __ne__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # type: ignore[misc] + @overload + def __ne__( # type: ignore[misc] self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... + @overload + def __ne__(self, other: object) -> Literal[True]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __mod__(self, other: timedelta) -> Timedelta: ... @@ -342,7 +352,7 @@ class Timedelta(timedelta): self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... @overload - def __le__(self, other: TimedeltaSeries) -> Series[bool]: ... + def __le__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __lt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @@ -351,7 +361,7 @@ class Timedelta(timedelta): self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... @overload - def __lt__(self, other: TimedeltaSeries) -> Series[bool]: ... + def __lt__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __ge__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @@ -360,7 +370,7 @@ class Timedelta(timedelta): self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... @overload - def __ge__(self, other: TimedeltaSeries) -> Series[bool]: ... + def __ge__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # Override due to more types supported than dt.timedelta @overload # type: ignore[override] def __gt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc] @@ -369,7 +379,7 @@ class Timedelta(timedelta): self, other: TimedeltaIndex | npt.NDArray[np.timedelta64] ) -> npt.NDArray[np.bool_]: ... @overload - def __gt__(self, other: TimedeltaSeries) -> Series[bool]: ... + def __gt__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... def isoformat(self) -> str: ... def to_numpy(self) -> np.timedelta64: ... @property diff --git a/tests/test_scalars.py b/tests/test_scalars.py index b8e3588bc..361d037b6 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -4,6 +4,7 @@ from typing import ( TYPE_CHECKING, Any, + Literal, cast, ) @@ -233,8 +234,9 @@ def test_timedelta_add_sub() -> None: # https://github.com/microsoft/pyright/issues/4088 check( assert_type( - as_dt_timedelta + td, pd.Timedelta - ), # pyright: ignore[reportGeneralTypeIssues] + as_dt_timedelta + td, + pd.Timedelta, # pyright: ignore[reportGeneralTypeIssues] + ), pd.Timedelta, ) check(assert_type(as_timedelta64 + td, pd.Timedelta), pd.Timedelta) @@ -321,8 +323,8 @@ def test_timedelta_mul_div() -> None: md_float = 3.5 md_ndarray_intp = np_intp_arr md_ndarray_float = np_float_arr - mp_series_int = pd.Series([1, 2, 3]) - md_series_float = pd.Series([1.2, 2.2, 3.4]) + mp_series_int = pd.Series([1, 2, 3], dtype=int) + md_series_float = pd.Series([1.2, 2.2, 3.4], dtype=float) md_int64_index = i_idx md_float_index = f_idx @@ -340,14 +342,15 @@ def test_timedelta_mul_div() -> None: # pyright is wrong here ndarray.__mul__(Timedelta0 is NotImplemented check( assert_type( - md_ndarray_intp * td, np.ndarray - ), # pyright: ignore[reportGeneralTypeIssues] + md_ndarray_intp * td, np.ndarray # pyright: ignore[reportGeneralTypeIssues] + ), np.ndarray, ) check( assert_type( - md_ndarray_float * td, np.ndarray - ), # pyright: ignore[reportGeneralTypeIssues] + md_ndarray_float * td, + np.ndarray, # pyright: ignore[reportGeneralTypeIssues] + ), np.ndarray, ) check(assert_type(mp_series_int * td, TimedeltaSeries), pd.Series) @@ -454,7 +457,7 @@ def test_timedelta_cmp() -> None: c_timedelta64 = np.timedelta64(1, "D") c_ndarray_td64 = ndarray_td64 c_timedelta_index = pd.TimedeltaIndex([1, 2, 3], unit="D") - c_timedelta_series = pd.Series([1, 2, 3], dtype="timedelta64[D]") + c_timedelta_series = pd.Series(pd.TimedeltaIndex([1, 2, 3])) check(assert_type(td < c_timedelta, bool), bool) check(assert_type(td < c_dt_timedelta, bool), bool) @@ -665,22 +668,14 @@ def test_timedelta_cmp() -> None: ) assert (eq_s != ne_s).all() - eq = check(assert_type(td == 1, bool), bool) - ne = check(assert_type(td != 1, bool), bool) + eq = check(assert_type(td == 1, Literal[False]), bool) + ne = check(assert_type(td != 1, Literal[True]), bool) assert eq != ne - eq = check(assert_type(td == (3 + 2j), bool), bool) - ne = check(assert_type(td != (3 + 2j), bool), bool) + eq = check(assert_type(td == (3 + 2j), Literal[False]), bool) + ne = check(assert_type(td != (3 + 2j), Literal[True]), bool) assert eq != ne - eq_s = check( - assert_type(td == pd.Series([1, 2, 3]), "pd.Series[bool]"), pd.Series, bool - ) - ne_s = check( - assert_type(td != pd.Series([1, 2, 3]), "pd.Series[bool]"), pd.Series, bool - ) - assert (eq_s != ne_s).all() - def test_period_construction() -> None: p = pd.Period("2012-1-1", freq="D") From e755b84d699b725944afe4f487d1308bb8b387c4 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Tue, 25 Oct 2022 07:58:17 +0100 Subject: [PATCH 07/10] Final changes to timedelta --- tests/test_scalars.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 361d037b6..3746f633f 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -234,8 +234,8 @@ def test_timedelta_add_sub() -> None: # https://github.com/microsoft/pyright/issues/4088 check( assert_type( - as_dt_timedelta + td, - pd.Timedelta, # pyright: ignore[reportGeneralTypeIssues] + as_dt_timedelta + td, # pyright: ignore[reportGeneralTypeIssues] + pd.Timedelta, ), pd.Timedelta, ) @@ -327,6 +327,7 @@ def test_timedelta_mul_div() -> None: md_series_float = pd.Series([1.2, 2.2, 3.4], dtype=float) md_int64_index = i_idx md_float_index = f_idx + md_timedelta_series = pd.Series(pd.timedelta_range("1 day", periods=3)) check(assert_type(td * md_int, pd.Timedelta), pd.Timedelta) check(assert_type(td * md_float, pd.Timedelta), pd.Timedelta) @@ -348,8 +349,8 @@ def test_timedelta_mul_div() -> None: ) check( assert_type( - md_ndarray_float * td, - np.ndarray, # pyright: ignore[reportGeneralTypeIssues] + md_ndarray_float * td, # pyright: ignore[reportGeneralTypeIssues] + np.ndarray, ), np.ndarray, ) @@ -368,6 +369,7 @@ def test_timedelta_mul_div() -> None: check(assert_type(td // md_series_float, TimedeltaSeries), pd.Series) check(assert_type(td // md_int64_index, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td // md_float_index, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td // md_timedelta_series, "pd.Series[int]"), pd.Series, int) check(assert_type(pd.NaT // td, float), float) # Note: None of the reverse floordiv work @@ -392,6 +394,7 @@ def test_timedelta_mul_div() -> None: check(assert_type(td / md_series_float, TimedeltaSeries), pd.Series) check(assert_type(td / md_int64_index, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td / md_float_index, pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(td / md_timedelta_series, "pd.Series[float]"), pd.Series, float) check(assert_type(pd.NaT / td, float), float) # Note: None of the reverse truediv work From 15be2c696b7f0870b0e50c6c8c835744636f1664 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Tue, 25 Oct 2022 11:01:47 +0100 Subject: [PATCH 08/10] Final cleanup --- tests/test_scalars.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 3746f633f..ec05ce6d3 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -229,8 +229,7 @@ def test_timedelta_add_sub() -> None: check(assert_type(as_datetime + td, dt.datetime), dt.datetime) check(assert_type(as_date + td, dt.date), dt.date) check(assert_type(as_datetime64 + td, pd.Timestamp), pd.Timestamp) - # pyright is fails here because as6 + td calls td.__radd__(as6), - # not timedelta.__add__ + # pyright can't know that as6 + td calls td.__radd__(as6), not timedelta.__sub__ # https://github.com/microsoft/pyright/issues/4088 check( assert_type( @@ -274,8 +273,8 @@ def test_timedelta_add_sub() -> None: check(assert_type(as_datetime - td, dt.datetime), dt.datetime) check(assert_type(as_date - td, dt.date), dt.date) check(assert_type(as_datetime64 - td, pd.Timestamp), pd.Timestamp) - # pyright is wrong here because as6 + td calls td.__rsub__(as6), - # not timedelta.__sub__ + # pyright can't know that as6 + td calls td.__rsub__(as6), not timedelta.__sub__ + # https://github.com/microsoft/pyright/issues/4088 check( assert_type( as_dt_timedelta - td, # pyright: ignore[reportGeneralTypeIssues] @@ -340,7 +339,7 @@ def test_timedelta_mul_div() -> None: check(assert_type(md_int * td, pd.Timedelta), pd.Timedelta) check(assert_type(md_float * td, pd.Timedelta), pd.Timedelta) - # pyright is wrong here ndarray.__mul__(Timedelta0 is NotImplemented + # pyright is wrong here ndarray.__mul__(Timedelta) is NotImplemented check( assert_type( md_ndarray_intp * td, np.ndarray # pyright: ignore[reportGeneralTypeIssues] @@ -443,7 +442,7 @@ def test_timedelta_mod_abs_unary() -> None: pd.TimedeltaIndex, ) - # mypy and pyright reports dt.timedelta, even though __abs__ returns Timedelta + # mypy and pyright report dt.timedelta, even though __abs__ returns Timedelta check(assert_type(abs(td), dt.timedelta), pd.Timedelta) check(assert_type(td.__abs__(), pd.Timedelta), pd.Timedelta) check(assert_type(-td, pd.Timedelta), pd.Timedelta) From bbc9cfeed5980cccb496c903ff6411cc103e674f Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Thu, 3 Nov 2022 17:51:54 +0000 Subject: [PATCH 09/10] CLN: Final clean --- pandas-stubs/_libs/tslibs/timedeltas.pyi | 4 +- pandas-stubs/core/algorithms.pyi | 1 - pandas-stubs/core/arrays/arrow/dtype.pyi | 1 - pandas-stubs/core/arrays/numpy_.pyi | 1 - pandas-stubs/core/dtypes/common.pyi | 1 - pandas-stubs/core/frame.pyi | 1 - pandas-stubs/core/groupby/generic.pyi | 1 - pandas-stubs/core/indexes/interval.pyi | 1 - pandas-stubs/core/indexes/numeric.pyi | 1 - pandas-stubs/core/ops/dispatch.pyi | 2 - pandas-stubs/core/series.pyi | 8 +- pandas-stubs/core/tools/datetimes.pyi | 1 - pandas-stubs/core/tools/timedeltas.pyi | 1 - pandas-stubs/core/window/expanding.pyi | 1 - pandas-stubs/core/window/rolling.pyi | 1 - pandas-stubs/io/clipboards.pyi | 1 - pandas-stubs/io/formats/format.pyi | 2 - pandas-stubs/io/parsers/readers.pyi | 1 - pandas-stubs/io/pytables.pyi | 1 - pandas-stubs/io/sas/sas7bdat.pyi | 2 - pandas-stubs/io/stata.pyi | 3 - pyproject.toml | 2 +- tests/test_scalars.py | 146 ++++++++++++++++------- 23 files changed, 110 insertions(+), 74 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/pandas-stubs/_libs/tslibs/timedeltas.pyi index 767d7b4a9..942e9ae20 100644 --- a/pandas-stubs/_libs/tslibs/timedeltas.pyi +++ b/pandas-stubs/_libs/tslibs/timedeltas.pyi @@ -231,7 +231,9 @@ class Timedelta(timedelta): @overload # type: ignore[override] def __mul__(self, other: float) -> Timedelta: ... @overload - def __mul__(self, other: np.ndarray) -> np.ndarray: ... + def __mul__( + self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating] + ) -> npt.NDArray[np.timedelta64]: ... @overload def __mul__(self, other: Series[int]) -> TimedeltaSeries: ... @overload diff --git a/pandas-stubs/core/algorithms.pyi b/pandas-stubs/core/algorithms.pyi index 4ac2621a1..4582c3184 100644 --- a/pandas-stubs/core/algorithms.pyi +++ b/pandas-stubs/core/algorithms.pyi @@ -4,7 +4,6 @@ from typing import ( ) import numpy as np -import pandas as pd from pandas import ( Categorical, CategoricalIndex, diff --git a/pandas-stubs/core/arrays/arrow/dtype.pyi b/pandas-stubs/core/arrays/arrow/dtype.pyi index dc27fff7c..041012a7a 100644 --- a/pandas-stubs/core/arrays/arrow/dtype.pyi +++ b/pandas-stubs/core/arrays/arrow/dtype.pyi @@ -1,4 +1,3 @@ -import numpy as np import pyarrow as pa from pandas.core.dtypes.base import StorageExtensionDtype diff --git a/pandas-stubs/core/arrays/numpy_.pyi b/pandas-stubs/core/arrays/numpy_.pyi index f44e5a6af..535b1905d 100644 --- a/pandas-stubs/core/arrays/numpy_.pyi +++ b/pandas-stubs/core/arrays/numpy_.pyi @@ -1,4 +1,3 @@ -import numpy as np from numpy.lib.mixins import NDArrayOperatorsMixin from pandas.core.arrays.base import ( ExtensionArray, diff --git a/pandas-stubs/core/dtypes/common.pyi b/pandas-stubs/core/dtypes/common.pyi index 3da5a0091..3156f1508 100644 --- a/pandas-stubs/core/dtypes/common.pyi +++ b/pandas-stubs/core/dtypes/common.pyi @@ -1,6 +1,5 @@ from typing import Union -import numpy as np import pandas as pd from typing_extensions import TypeAlias diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 7d49cf35b..d82f70c33 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -110,7 +110,6 @@ from pandas._typing import ( num, ) -from pandas.io.formats import format as fmt from pandas.io.formats.style import Styler from pandas.plotting import PlotAccessor diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi index 0cc15a2b9..3dff35b2c 100644 --- a/pandas-stubs/core/groupby/generic.pyi +++ b/pandas-stubs/core/groupby/generic.pyi @@ -15,7 +15,6 @@ from matplotlib.axes import ( Axes as PlotAxes, SubplotBase as AxesSubplot, ) -import numpy as np from pandas.core.frame import DataFrame from pandas.core.generic import NDFrame from pandas.core.groupby.groupby import ( # , get_groupby as get_groupby diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index 339f99053..dd5ef8510 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -1,4 +1,3 @@ -import datetime as dt from typing import ( Any, Hashable, diff --git a/pandas-stubs/core/indexes/numeric.pyi b/pandas-stubs/core/indexes/numeric.pyi index 5e927327e..95715a83a 100644 --- a/pandas-stubs/core/indexes/numeric.pyi +++ b/pandas-stubs/core/indexes/numeric.pyi @@ -1,6 +1,5 @@ from typing import Iterable -import numpy as np from pandas.core.indexes.base import Index from pandas._typing import np_ndarray_int64 diff --git a/pandas-stubs/core/ops/dispatch.pyi b/pandas-stubs/core/ops/dispatch.pyi index 0a69d985a..b283ba29e 100644 --- a/pandas-stubs/core/ops/dispatch.pyi +++ b/pandas-stubs/core/ops/dispatch.pyi @@ -1,5 +1,3 @@ -import numpy as np - from pandas.core.dtypes.generic import ABCSeries def should_extension_dispatch(left: ABCSeries, right) -> bool: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 6139dfe61..e368c8786 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1303,8 +1303,12 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): ) -> TimedeltaSeries: ... @overload def __sub__( - self, other: Timedelta | TimedeltaSeries | TimedeltaIndex - ) -> TimestampSeries | TimedeltaSeries: ... + self: Series[Timestamp], other: Timedelta | TimedeltaSeries | TimedeltaIndex + ) -> TimestampSeries: ... + @overload + def __sub__( + self: Series[Timedelta], other: Timedelta | TimedeltaSeries | TimedeltaIndex + ) -> TimedeltaSeries: ... @overload def __sub__(self, other: num | _ListLike | Series) -> Series: ... @overload diff --git a/pandas-stubs/core/tools/datetimes.pyi b/pandas-stubs/core/tools/datetimes.pyi index fa32eceff..6cc504f7d 100644 --- a/pandas-stubs/core/tools/datetimes.pyi +++ b/pandas-stubs/core/tools/datetimes.pyi @@ -10,7 +10,6 @@ from typing import ( ) import numpy as np -import pandas as pd from pandas import ( Index, Timestamp, diff --git a/pandas-stubs/core/tools/timedeltas.pyi b/pandas-stubs/core/tools/timedeltas.pyi index 5f6817b8f..ed1818348 100644 --- a/pandas-stubs/core/tools/timedeltas.pyi +++ b/pandas-stubs/core/tools/timedeltas.pyi @@ -4,7 +4,6 @@ from typing import ( overload, ) -import pandas as pd from pandas import Index from pandas.core.indexes.timedeltas import TimedeltaIndex from pandas.core.series import ( diff --git a/pandas-stubs/core/window/expanding.pyi b/pandas-stubs/core/window/expanding.pyi index 1c3b164ff..04df37896 100644 --- a/pandas-stubs/core/window/expanding.pyi +++ b/pandas-stubs/core/window/expanding.pyi @@ -4,7 +4,6 @@ from typing import ( overload, ) -import numpy as np from pandas import ( DataFrame, Series, diff --git a/pandas-stubs/core/window/rolling.pyi b/pandas-stubs/core/window/rolling.pyi index de3eb8357..9830b01d3 100644 --- a/pandas-stubs/core/window/rolling.pyi +++ b/pandas-stubs/core/window/rolling.pyi @@ -5,7 +5,6 @@ from typing import ( overload, ) -import numpy as np from pandas import ( DataFrame, Series, diff --git a/pandas-stubs/io/clipboards.pyi b/pandas-stubs/io/clipboards.pyi index a57c5ecab..cb8449902 100644 --- a/pandas-stubs/io/clipboards.pyi +++ b/pandas-stubs/io/clipboards.pyi @@ -8,7 +8,6 @@ from typing import ( overload, ) -import numpy as np from pandas.core.frame import DataFrame from pandas.core.indexes.base import Index from pandas.core.series import Series diff --git a/pandas-stubs/io/formats/format.pyi b/pandas-stubs/io/formats/format.pyi index 0305a2886..c26fbcf2f 100644 --- a/pandas-stubs/io/formats/format.pyi +++ b/pandas-stubs/io/formats/format.pyi @@ -1,5 +1,3 @@ -import numpy as np - class EngFormatter: ENG_PREFIXES = ... accuracy = ... diff --git a/pandas-stubs/io/parsers/readers.pyi b/pandas-stubs/io/parsers/readers.pyi index 14db44693..cf2d7d3f1 100644 --- a/pandas-stubs/io/parsers/readers.pyi +++ b/pandas-stubs/io/parsers/readers.pyi @@ -16,7 +16,6 @@ from pandas.core.frame import DataFrame from pandas.core.indexes.base import Index from pandas.core.series import Series -import pandas._libs.lib as lib from pandas._typing import ( CompressionOptions, CSVEngine, diff --git a/pandas-stubs/io/pytables.pyi b/pandas-stubs/io/pytables.pyi index 4dcaea325..41719e389 100644 --- a/pandas-stubs/io/pytables.pyi +++ b/pandas-stubs/io/pytables.pyi @@ -7,7 +7,6 @@ from typing import ( overload, ) -import numpy as np from pandas import ( DataFrame, Series, diff --git a/pandas-stubs/io/sas/sas7bdat.pyi b/pandas-stubs/io/sas/sas7bdat.pyi index fdb0a556a..1db480959 100644 --- a/pandas-stubs/io/sas/sas7bdat.pyi +++ b/pandas-stubs/io/sas/sas7bdat.pyi @@ -1,5 +1,3 @@ -import numpy as np -import pandas as pd from pandas import DataFrame from pandas.io.sas.sasreader import ReaderBase diff --git a/pandas-stubs/io/stata.pyi b/pandas-stubs/io/stata.pyi index 86c8b0825..aca173e87 100644 --- a/pandas-stubs/io/stata.pyi +++ b/pandas-stubs/io/stata.pyi @@ -8,9 +8,6 @@ from typing import ( overload, ) -import numpy as np -import numpy.typing as npt -import pandas as pd from pandas.core.frame import DataFrame from pandas._typing import ( diff --git a/pyproject.toml b/pyproject.toml index 457323444..3a68d9e9d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ types-pytz = ">= 2022.1.1" mypy = "==0.971" pyarrow = ">=9.0.0" pytest = ">=7.1.2" -pyright = ">=1.1.266" +pyright = ">=1.1.278" poethepoet = "0.16.0" loguru = ">=0.6.0" pandas = "1.5.1" diff --git a/tests/test_scalars.py b/tests/test_scalars.py index ec05ce6d3..4b1baea86 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -205,7 +205,9 @@ def test_timedelta_add_sub() -> None: check(assert_type(td + as_dt_timedelta, pd.Timedelta), pd.Timedelta) check(assert_type(td + as_timedelta64, pd.Timedelta), pd.Timedelta) check(assert_type(td + as_timedelta_index, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(td + as_timedelta_series, TimedeltaSeries), pd.Series) + check( + assert_type(td + as_timedelta_series, TimedeltaSeries), pd.Series, pd.Timedelta + ) check(assert_type(td + as_period_index, pd.PeriodIndex), pd.PeriodIndex) check(assert_type(td + as_datetime_index, pd.DatetimeIndex), pd.DatetimeIndex) check( @@ -213,12 +215,8 @@ def test_timedelta_add_sub() -> None: np.ndarray, np.timedelta64, ) - # pyright has trouble with timedelta64 and datetime64 check( - assert_type( - td + as_ndarray_dt64, # pyright: ignore[reportGeneralTypeIssues] - npt.NDArray[np.datetime64], - ), + assert_type(td + as_ndarray_dt64, npt.NDArray[np.datetime64]), np.ndarray, np.datetime64, ) @@ -240,7 +238,9 @@ def test_timedelta_add_sub() -> None: ) check(assert_type(as_timedelta64 + td, pd.Timedelta), pd.Timedelta) check(assert_type(as_timedelta_index + td, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(as_timedelta_series + td, TimedeltaSeries), pd.Series) + check( + assert_type(as_timedelta_series + td, TimedeltaSeries), pd.Series, pd.Timedelta + ) check(assert_type(as_period_index + td, pd.PeriodIndex), pd.PeriodIndex) check(assert_type(as_datetime_index + td, pd.DatetimeIndex), pd.DatetimeIndex) # pyright is fails here because ndarray.__add__(Timedelta) is NotImplemented @@ -256,12 +256,25 @@ def test_timedelta_add_sub() -> None: # sub is not symmetric with dates. In general date_like - timedelta is # sensible, while timedelta - date_like is not - # TypeError: as1, as2, as3, as4, as5, as10, as11, as13 + # TypeError: as_period, as_timestamp, as_datetime, as_date, as_datetime64, + # as_period_index, as_datetime_index, as_ndarray_dt64 + if TYPE_CHECKING_INVALID_USAGE: + td - as_period # type: ignore[operator] + td - as_timestamp # type: ignore[operator] + td - as_datetime # type: ignore[operator] + td - as_date # type: ignore[operator] + td - as_datetime64 # type: ignore[operator] + td - as_period_index # type: ignore[operator] + td - as_datetime_index # type: ignore[operator] + td - as_ndarray_dt64 # type: ignore[operator] + check(assert_type(td - td, pd.Timedelta), pd.Timedelta) check(assert_type(td - as_dt_timedelta, pd.Timedelta), pd.Timedelta) check(assert_type(td - as_timedelta64, pd.Timedelta), pd.Timedelta) check(assert_type(td - as_timedelta_index, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(td - as_timedelta_series, TimedeltaSeries), pd.Series) + check( + assert_type(td - as_timedelta_series, TimedeltaSeries), pd.Series, pd.Timedelta + ) check( assert_type(td - as_ndarray_td64, npt.NDArray[np.timedelta64]), np.ndarray, @@ -273,7 +286,8 @@ def test_timedelta_add_sub() -> None: check(assert_type(as_datetime - td, dt.datetime), dt.datetime) check(assert_type(as_date - td, dt.date), dt.date) check(assert_type(as_datetime64 - td, pd.Timestamp), pd.Timestamp) - # pyright can't know that as6 + td calls td.__rsub__(as6), not timedelta.__sub__ + # pyright can't know that as_dt_timedelta - td calls td.__rsub__(as_dt_timedelta), + # not as_dt_timedelta.__sub__ # https://github.com/microsoft/pyright/issues/4088 check( assert_type( @@ -284,22 +298,22 @@ def test_timedelta_add_sub() -> None: ) check(assert_type(as_timedelta64 - td, pd.Timedelta), pd.Timedelta) check(assert_type(as_timedelta_index - td, pd.TimedeltaIndex), pd.TimedeltaIndex) - check(assert_type(as_timedelta_series - td, TimedeltaSeries), pd.Series) + check( + assert_type(as_timedelta_series - td, TimedeltaSeries), pd.Series, pd.Timedelta + ) check(assert_type(as_period_index - td, pd.PeriodIndex), pd.PeriodIndex) check(assert_type(as_datetime_index - td, pd.DatetimeIndex), pd.DatetimeIndex) - # pyright is wrong here because ndarray.__sub__(Timedelta) is NotImplemented check( assert_type( - as_ndarray_td64 - td, # pyright: ignore[reportGeneralTypeIssues] + as_ndarray_td64 - td, npt.NDArray[np.timedelta64], ), np.ndarray, np.timedelta64, ) - # pyright is wrong here because ndarray.__sub__(Timedelta) is NotImplemented check( assert_type( - as_ndarray_dt64 - td, # pyright: ignore[reportGeneralTypeIssues] + as_ndarray_dt64 - td, npt.NDArray[np.datetime64], ), np.ndarray, @@ -330,10 +344,18 @@ def test_timedelta_mul_div() -> None: check(assert_type(td * md_int, pd.Timedelta), pd.Timedelta) check(assert_type(td * md_float, pd.Timedelta), pd.Timedelta) - check(assert_type(td * md_ndarray_intp, np.ndarray), np.ndarray) - check(assert_type(td * md_ndarray_float, np.ndarray), np.ndarray) - check(assert_type(td * mp_series_int, TimedeltaSeries), pd.Series) - check(assert_type(td * md_series_float, TimedeltaSeries), pd.Series) + check( + assert_type(td * md_ndarray_intp, npt.NDArray[np.timedelta64]), + np.ndarray, + np.timedelta64, + ) + check( + assert_type(td * md_ndarray_float, npt.NDArray[np.timedelta64]), + np.ndarray, + np.timedelta64, + ) + check(assert_type(td * mp_series_int, TimedeltaSeries), pd.Series, pd.Timedelta) + check(assert_type(td * md_series_float, TimedeltaSeries), pd.Series, pd.Timedelta) check(assert_type(td * md_int64_index, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td * md_float_index, pd.TimedeltaIndex), pd.TimedeltaIndex) @@ -353,8 +375,8 @@ def test_timedelta_mul_div() -> None: ), np.ndarray, ) - check(assert_type(mp_series_int * td, TimedeltaSeries), pd.Series) - check(assert_type(md_series_float * td, TimedeltaSeries), pd.Series) + check(assert_type(mp_series_int * td, TimedeltaSeries), pd.Series, pd.Timedelta) + check(assert_type(md_series_float * td, TimedeltaSeries), pd.Series, pd.Timedelta) check(assert_type(md_int64_index * td, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(md_float_index * td, pd.TimedeltaIndex), pd.TimedeltaIndex) @@ -362,17 +384,26 @@ def test_timedelta_mul_div() -> None: check(assert_type(td // pd.NaT, float), float) check(assert_type(td // md_int, pd.Timedelta), pd.Timedelta) check(assert_type(td // md_float, pd.Timedelta), pd.Timedelta) - check(assert_type(td // md_ndarray_intp, npt.NDArray[np.timedelta64]), np.ndarray) - check(assert_type(td // md_ndarray_float, npt.NDArray[np.timedelta64]), np.ndarray) - check(assert_type(td // mp_series_int, TimedeltaSeries), pd.Series) - check(assert_type(td // md_series_float, TimedeltaSeries), pd.Series) + check( + assert_type(td // md_ndarray_intp, npt.NDArray[np.timedelta64]), + np.ndarray, + np.timedelta64, + ) + check( + assert_type(td // md_ndarray_float, npt.NDArray[np.timedelta64]), + np.ndarray, + np.timedelta64, + ) + check(assert_type(td // mp_series_int, TimedeltaSeries), pd.Series, pd.Timedelta) + check(assert_type(td // md_series_float, TimedeltaSeries), pd.Series, pd.Timedelta) check(assert_type(td // md_int64_index, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td // md_float_index, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td // md_timedelta_series, "pd.Series[int]"), pd.Series, int) check(assert_type(pd.NaT // td, float), float) # Note: None of the reverse floordiv work - # TypeError: md1, md2, md3, md4, md5, md6, md7, md8 + # TypeError: md_int, md_float, md_ndarray_intp, md_ndarray_float, mp_series_int, + # mp_series_float, md_int64_index, md_float_index if TYPE_CHECKING_INVALID_USAGE: md_int // td # type: ignore[operator] md_float // td # type: ignore[operator] @@ -387,17 +418,26 @@ def test_timedelta_mul_div() -> None: check(assert_type(td / pd.NaT, float), float) check(assert_type(td / md_int, pd.Timedelta), pd.Timedelta) check(assert_type(td / md_float, pd.Timedelta), pd.Timedelta) - check(assert_type(td / md_ndarray_intp, npt.NDArray[np.timedelta64]), np.ndarray) - check(assert_type(td / md_ndarray_float, npt.NDArray[np.timedelta64]), np.ndarray) - check(assert_type(td / mp_series_int, TimedeltaSeries), pd.Series) - check(assert_type(td / md_series_float, TimedeltaSeries), pd.Series) + check( + assert_type(td / md_ndarray_intp, npt.NDArray[np.timedelta64]), + np.ndarray, + np.timedelta64, + ) + check( + assert_type(td / md_ndarray_float, npt.NDArray[np.timedelta64]), + np.ndarray, + np.timedelta64, + ) + check(assert_type(td / mp_series_int, TimedeltaSeries), pd.Series, pd.Timedelta) + check(assert_type(td / md_series_float, TimedeltaSeries), pd.Series, pd.Timedelta) check(assert_type(td / md_int64_index, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td / md_float_index, pd.TimedeltaIndex), pd.TimedeltaIndex) check(assert_type(td / md_timedelta_series, "pd.Series[float]"), pd.Series, float) check(assert_type(pd.NaT / td, float), float) # Note: None of the reverse truediv work - # TypeError: md1, md2, md3, md4, md5, md6, md7, md8 + # TypeError: md_int, md_float, md_ndarray_intp, md_ndarray_float, mp_series_int, + # mp_series_float, md_int64_index, md_float_index if TYPE_CHECKING_INVALID_USAGE: md_int / td # type: ignore[operator] md_float / td # type: ignore[operator] @@ -405,8 +445,8 @@ def test_timedelta_mul_div() -> None: md_ndarray_float / td # type: ignore[operator] # TODO: Series.__truediv__ says it supports Timedelta # it does not, in general, except for TimedeltaSeries - # md5 / td # type: ignore[operator] - # md6 / td # type: ignore[operator] + # mp_series_int / td # type: ignore[operator] + # mp_series_float / td # type: ignore[operator] md_int64_index / td # type: ignore[operator] md_float_index / td # type: ignore[operator] @@ -433,8 +473,8 @@ def test_timedelta_mod_abs_unary() -> None: ) int_series = pd.Series([1, 2, 3], dtype=int) float_series = pd.Series([1.2, 2.2, 3.4], dtype=float) - check(assert_type(td % int_series, TimedeltaSeries), pd.Series) - check(assert_type(td % float_series, TimedeltaSeries), pd.Series) + check(assert_type(td % int_series, TimedeltaSeries), pd.Series, pd.Timedelta) + check(assert_type(td % float_series, TimedeltaSeries), pd.Series, pd.Timedelta) check(assert_type(td % i_idx, pd.TimedeltaIndex), pd.TimedeltaIndex) check( @@ -442,8 +482,8 @@ def test_timedelta_mod_abs_unary() -> None: pd.TimedeltaIndex, ) - # mypy and pyright report dt.timedelta, even though __abs__ returns Timedelta - check(assert_type(abs(td), dt.timedelta), pd.Timedelta) + # mypy reports dt.timedelta, even though __abs__ returns Timedelta + check(assert_type(abs(td), pd.Timedelta), pd.Timedelta) # type: ignore[assert-type] 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) @@ -642,6 +682,28 @@ def test_timedelta_cmp() -> None: ) assert (eq_s != ne_s).all() + eq = check(assert_type(td == 1, Literal[False]), bool) + ne = check(assert_type(td != 1, Literal[True]), bool) + assert eq != ne + + eq = check(assert_type(td == (3 + 2j), Literal[False]), bool) + ne = check(assert_type(td != (3 + 2j), Literal[True]), bool) + assert eq != ne + + +def test_timedelta_cmp_rhs() -> None: + # Test that check eq and ne when Timedelta is the RHS argument + # that use the __eq__ and __ne__ methods of the LHS + td = pd.Timedelta("1 day") + ndarray_td64: npt.NDArray[np.timedelta64] = np.array( + [1, 2, 3], dtype="timedelta64[D]" + ) + c_dt_timedelta = dt.timedelta(days=1) + c_timedelta64 = np.timedelta64(1, "D") + c_ndarray_td64 = ndarray_td64 + c_timedelta_index = pd.TimedeltaIndex([1, 2, 3], unit="D") + c_timedelta_series = pd.Series(pd.TimedeltaIndex([1, 2, 3])) + eq = check(assert_type(c_dt_timedelta == td, bool), bool) ne = check(assert_type(c_dt_timedelta != td, bool), bool) assert eq != ne @@ -670,14 +732,6 @@ def test_timedelta_cmp() -> None: ) assert (eq_s != ne_s).all() - eq = check(assert_type(td == 1, Literal[False]), bool) - ne = check(assert_type(td != 1, Literal[True]), bool) - assert eq != ne - - eq = check(assert_type(td == (3 + 2j), Literal[False]), bool) - ne = check(assert_type(td != (3 + 2j), Literal[True]), bool) - assert eq != ne - def test_period_construction() -> None: p = pd.Period("2012-1-1", freq="D") @@ -752,7 +806,7 @@ def test_period_add_subtract() -> None: as_period = pd.Period("2012-1-1", freq="D") scale = 24 * 60 * 60 * 10**9 as_td_series = pd.Series(pd.timedelta_range(scale, scale, freq="D")) - check(assert_type(as_td_series, TimedeltaSeries), pd.Series) + check(assert_type(as_td_series, TimedeltaSeries), pd.Series, pd.Timedelta) as_period_series = pd.Series(as_period_index) check(assert_type(as_period_series, PeriodSeries), pd.Series) as_timedelta_idx = pd.timedelta_range(scale, scale, freq="D") From e877b6f6907ffcd61480ccbf9f9ed6be836a2741 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Thu, 3 Nov 2022 17:55:32 +0000 Subject: [PATCH 10/10] CLN: Final clean --- tests/test_scalars.py | 19 +++---------------- 1 file changed, 3 insertions(+), 16 deletions(-) diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 4b1baea86..0c8d03a81 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -243,10 +243,9 @@ def test_timedelta_add_sub() -> None: ) check(assert_type(as_period_index + td, pd.PeriodIndex), pd.PeriodIndex) check(assert_type(as_datetime_index + td, pd.DatetimeIndex), pd.DatetimeIndex) - # pyright is fails here because ndarray.__add__(Timedelta) is NotImplemented check( assert_type( - as_ndarray_td64 + td, # pyright: ignore[reportGeneralTypeIssues] + as_ndarray_td64 + td, npt.NDArray[np.timedelta64], ), np.ndarray, @@ -361,20 +360,8 @@ def test_timedelta_mul_div() -> None: check(assert_type(md_int * td, pd.Timedelta), pd.Timedelta) check(assert_type(md_float * td, pd.Timedelta), pd.Timedelta) - # pyright is wrong here ndarray.__mul__(Timedelta) is NotImplemented - check( - assert_type( - md_ndarray_intp * td, np.ndarray # pyright: ignore[reportGeneralTypeIssues] - ), - np.ndarray, - ) - check( - assert_type( - md_ndarray_float * td, # pyright: ignore[reportGeneralTypeIssues] - np.ndarray, - ), - np.ndarray, - ) + check(assert_type(md_ndarray_intp * td, np.ndarray), np.ndarray) + check(assert_type(md_ndarray_float * td, np.ndarray), np.ndarray) check(assert_type(mp_series_int * td, TimedeltaSeries), pd.Series, pd.Timedelta) check(assert_type(md_series_float * td, TimedeltaSeries), pd.Series, pd.Timedelta) check(assert_type(md_int64_index * td, pd.TimedeltaIndex), pd.TimedeltaIndex)