From 5c79550afcd5b81f80ed1cdba0d26bb6597de6f6 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Fri, 18 Feb 2022 08:06:09 -0500 Subject: [PATCH 1/8] get types of Interval correct --- partial/pandas/_libs/interval.pyi | 105 ++++++++++++++++++++++-------- tests/pandas/test_interval.py | 7 ++ 2 files changed, 86 insertions(+), 26 deletions(-) diff --git a/partial/pandas/_libs/interval.pyi b/partial/pandas/_libs/interval.pyi index fb6cb7cf..33e9cd42 100644 --- a/partial/pandas/_libs/interval.pyi +++ b/partial/pandas/_libs/interval.pyi @@ -1,6 +1,6 @@ from __future__ import annotations import sys -from typing import Generic, overload, Union, Protocol, TypeVar +from typing import Generic, overload, Union, Protocol, TypeVar, Any from _typing import Timedelta, Timestamp import datetime @@ -9,12 +9,28 @@ if sys.version_info >= (3, 8): else: from typing_extensions import Literal -OrderableScalar = TypeVar("OrderableScalar", int, float) -OrderableTimes = TypeVar("OrderableTimes", datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta) -Orderable = TypeVar("Orderable", int, float, datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta) +OrderableScalarT = TypeVar("OrderableScalarT", int, float) +OrderableTimesT = TypeVar("OrderableTimesT", datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta) +OrderableDTimesT = TypeVar("OrderableDTimesT", datetime.date, datetime.datetime, datetime.timedelta) +OrderablePTimesT = TypeVar("OrderablePTimesT", Timestamp, Timedelta) +OrderableT = TypeVar("OrderableT", int, float, datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta) + +OrderableScalar = Union[int, float] +OrderableTimes = Union[datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta] +Orderable = Union[int, float, datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta] class IntervalMixinProtocol(Protocol): ... +class _LengthDescriptor: + @overload + def __get__(self, instance: Interval[float], owner: Any) -> float: ... + @overload + def __get__(self, instance: Interval[int], owner: Any) -> int: ... + @overload + def __get__(self, instance: Interval[OrderablePTimesT], owner: Any) -> Timedelta: ... + @overload + def __get__(self, instance: Interval[OrderableDTimesT], owner: Any) -> datetime.timedelta: ... + class IntervalMixin: @property def closed_left(self: IntervalMixinProtocol) -> bool: ... @@ -29,43 +45,80 @@ class IntervalMixin: @property def is_empty(self: IntervalMixinProtocol) -> bool: ... -class Interval(IntervalMixin, Generic[Orderable]): +class Interval(IntervalMixin, Generic[OrderableT]): @overload - def left(self: Interval[OrderableScalar]) -> OrderableScalar: ... + def left(self: Interval[OrderableScalarT]) -> OrderableScalar: ... @overload - def left(self: Interval[OrderableTimes]) -> OrderableTimes: ... + def left(self: Interval[OrderableTimesT]) -> OrderableTimes: ... @property @overload - def left(self: Interval[Orderable]) -> Orderable: ... + def left(self: Interval[OrderableT]) -> Orderable: ... @overload - def right(self: Interval[OrderableScalar]) -> OrderableScalar: ... + def right(self: Interval[OrderableScalarT]) -> OrderableScalar: ... @overload - def right(self: Interval[OrderableTimes]) -> OrderableTimes: ... + def right(self: Interval[OrderableTimesT]) -> OrderableTimes: ... @property @overload - def right(self: Interval[Orderable]) -> Orderable: ... + def right(self: Interval[OrderableT]) -> Orderable: ... @property def closed(self) -> str: ... - def __init__( - self, - left: Orderable, - right: Orderable, + @overload + def __new__( + cls, + left: Timestamp, + right: Timestamp, closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., - ) -> None: ... + ) -> Interval[Timestamp]: ... @overload - def length(self: Interval[OrderableScalar]) -> float: ... + def __new__( + cls, + left: Timedelta, + right: Timedelta, + closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., + ) -> Interval[Timedelta]: ... @overload - def length(self: Interval[OrderableTimes]) -> Timedelta: ... - @property + def __new__( + cls, + left: datetime.date, + right: datetime.date, + closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., + ) -> Interval[datetime.date]: ... + @overload + def __new__( + cls, + left: datetime.datetime, + right: datetime.datetime, + closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., + ) -> Interval[datetime.datetime]: ... + @overload + def __new__( + cls, + left: datetime.timedelta, + right: datetime.timedelta, + closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., + ) -> Interval[datetime.timedelta]: ... + @overload + def __new__( + cls, + left: int, + right: int, + closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., + ) -> Interval[int]: ... @overload - def length(self: Interval[Orderable]) -> Orderable: ... + def __new__( + cls, + left: float, + right: float, + closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., + ) -> Interval[float]: ... + length: _LengthDescriptor def __hash__(self) -> int: ... def __contains__(self, key: Orderable) -> bool: ... def __repr__(self) -> str: ... def __str__(self) -> str: ... - def __add__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ... - def __sub__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ... - def __mul__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ... - def __truediv__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ... - def __floordiv__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ... - def overlaps(self: Interval[Orderable], other: Interval[Orderable]) -> bool: ... + def __add__(self: Interval[OrderableT], y: float) -> Interval[OrderableT]: ... + def __sub__(self: Interval[OrderableT], y: float) -> Interval[OrderableT]: ... + def __mul__(self: Interval[OrderableT], y: float) -> Interval[OrderableT]: ... + def __truediv__(self: Interval[OrderableT], y: float) -> Interval[OrderableT]: ... + def __floordiv__(self: Interval[OrderableT], y: float) -> Interval[OrderableT]: ... + def overlaps(self: Interval[OrderableT], other: Interval[OrderableT]) -> bool: ... diff --git a/tests/pandas/test_interval.py b/tests/pandas/test_interval.py index 20a064ae..c637d065 100644 --- a/tests/pandas/test_interval.py +++ b/tests/pandas/test_interval.py @@ -28,4 +28,11 @@ def test_max_intervals() -> None: def test_interval_length() -> None: i1 = pd.Interval(pd.Timestamp("2000-01-01"), pd.Timestamp("2000-01-02"), closed="both") + reveal_type(i1.length, expected_string="Timedelta") i1.length.total_seconds() + + i2 = pd.Interval(10, 20) + reveal_type(i2.length, expected_type=int) + + i3 = pd.Interval(13.2, 19.5) + reveal_type(i3.length, expected_type=float) From e93d9ab123507e30707c3addc035aaf16ab13e5e Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Sun, 20 Feb 2022 11:50:38 -0500 Subject: [PATCH 2/8] remove datetime --- partial/pandas/_libs/interval.pyi | 36 +++++-------------------------- 1 file changed, 5 insertions(+), 31 deletions(-) diff --git a/partial/pandas/_libs/interval.pyi b/partial/pandas/_libs/interval.pyi index 33e9cd42..c35fdbb8 100644 --- a/partial/pandas/_libs/interval.pyi +++ b/partial/pandas/_libs/interval.pyi @@ -2,7 +2,6 @@ from __future__ import annotations import sys from typing import Generic, overload, Union, Protocol, TypeVar, Any from _typing import Timedelta, Timestamp -import datetime if sys.version_info >= (3, 8): from typing import Literal @@ -10,14 +9,12 @@ else: from typing_extensions import Literal OrderableScalarT = TypeVar("OrderableScalarT", int, float) -OrderableTimesT = TypeVar("OrderableTimesT", datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta) -OrderableDTimesT = TypeVar("OrderableDTimesT", datetime.date, datetime.datetime, datetime.timedelta) -OrderablePTimesT = TypeVar("OrderablePTimesT", Timestamp, Timedelta) -OrderableT = TypeVar("OrderableT", int, float, datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta) +OrderableTimesT = TypeVar("OrderableTimesT", Timestamp, Timedelta) +OrderableT = TypeVar("OrderableT", int, float, Timestamp, Timedelta) OrderableScalar = Union[int, float] -OrderableTimes = Union[datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta] -Orderable = Union[int, float, datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta] +OrderableTimes = Union[Timestamp, Timedelta] +Orderable = Union[int, float, Timestamp, Timedelta] class IntervalMixinProtocol(Protocol): ... @@ -27,9 +24,7 @@ class _LengthDescriptor: @overload def __get__(self, instance: Interval[int], owner: Any) -> int: ... @overload - def __get__(self, instance: Interval[OrderablePTimesT], owner: Any) -> Timedelta: ... - @overload - def __get__(self, instance: Interval[OrderableDTimesT], owner: Any) -> datetime.timedelta: ... + def __get__(self, instance: Interval[OrderableTimesT], owner: Any) -> Timedelta: ... class IntervalMixin: @property @@ -77,27 +72,6 @@ class Interval(IntervalMixin, Generic[OrderableT]): closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., ) -> Interval[Timedelta]: ... @overload - def __new__( - cls, - left: datetime.date, - right: datetime.date, - closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., - ) -> Interval[datetime.date]: ... - @overload - def __new__( - cls, - left: datetime.datetime, - right: datetime.datetime, - closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., - ) -> Interval[datetime.datetime]: ... - @overload - def __new__( - cls, - left: datetime.timedelta, - right: datetime.timedelta, - closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., - ) -> Interval[datetime.timedelta]: ... - @overload def __new__( cls, left: int, From 2ed0928bba30a45f3fe824e382bdea6da0489d6d Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Sun, 20 Feb 2022 20:01:17 -0500 Subject: [PATCH 3/8] fix operators. use timestamp and timedelta from pandas --- partial/pandas/_libs/__init__.pyi | 4 + partial/pandas/_libs/interval.pyi | 66 +++-- partial/pandas/_libs/tslibs/timedeltas.pyi | 140 +++++----- partial/pandas/_libs/tslibs/timestamps.pyi | 302 +++++++++++---------- tests/pandas/test_interval.py | 40 ++- 5 files changed, 321 insertions(+), 231 deletions(-) create mode 100644 partial/pandas/_libs/__init__.pyi diff --git a/partial/pandas/_libs/__init__.pyi b/partial/pandas/_libs/__init__.pyi new file mode 100644 index 00000000..cdea8462 --- /dev/null +++ b/partial/pandas/_libs/__init__.pyi @@ -0,0 +1,4 @@ +from .tslibs import NaT as NaT, NaTType as NaTType, OutOfBoundsDatetime as OutOfBoundsDatetime, Period as Period, Timedelta as Timedelta, Timestamp as Timestamp, iNaT as iNaT +from .interval import Interval as Interval + + diff --git a/partial/pandas/_libs/interval.pyi b/partial/pandas/_libs/interval.pyi index c35fdbb8..a8a969bb 100644 --- a/partial/pandas/_libs/interval.pyi +++ b/partial/pandas/_libs/interval.pyi @@ -12,10 +12,6 @@ OrderableScalarT = TypeVar("OrderableScalarT", int, float) OrderableTimesT = TypeVar("OrderableTimesT", Timestamp, Timedelta) OrderableT = TypeVar("OrderableT", int, float, Timestamp, Timedelta) -OrderableScalar = Union[int, float] -OrderableTimes = Union[Timestamp, Timedelta] -Orderable = Union[int, float, Timestamp, Timedelta] - class IntervalMixinProtocol(Protocol): ... class _LengthDescriptor: @@ -41,20 +37,10 @@ class IntervalMixin: def is_empty(self: IntervalMixinProtocol) -> bool: ... class Interval(IntervalMixin, Generic[OrderableT]): - @overload - def left(self: Interval[OrderableScalarT]) -> OrderableScalar: ... - @overload - def left(self: Interval[OrderableTimesT]) -> OrderableTimes: ... @property - @overload - def left(self: Interval[OrderableT]) -> Orderable: ... - @overload - def right(self: Interval[OrderableScalarT]) -> OrderableScalar: ... - @overload - def right(self: Interval[OrderableTimesT]) -> OrderableTimes: ... + def left(self: Interval[OrderableT]) -> OrderableT: ... @property - @overload - def right(self: Interval[OrderableT]) -> Orderable: ... + def right(self: Interval[OrderableT]) -> OrderableT: ... @property def closed(self) -> str: ... @overload @@ -87,12 +73,48 @@ class Interval(IntervalMixin, Generic[OrderableT]): ) -> Interval[float]: ... length: _LengthDescriptor def __hash__(self) -> int: ... - def __contains__(self, key: Orderable) -> bool: ... + @overload + def __contains__(self: Interval[Timestamp], key: Timestamp) -> bool: ... + @overload + def __contains__(self: Interval[Timedelta], key: Timedelta) -> bool: ... + @overload + def __contains__(self: Interval[int], key: Union[int, float]) -> bool: ... + @overload + def __contains__(self: Interval[float], key: Union[int, float]) -> bool: ... def __repr__(self) -> str: ... def __str__(self) -> str: ... - def __add__(self: Interval[OrderableT], y: float) -> Interval[OrderableT]: ... - def __sub__(self: Interval[OrderableT], y: float) -> Interval[OrderableT]: ... - def __mul__(self: Interval[OrderableT], y: float) -> Interval[OrderableT]: ... - def __truediv__(self: Interval[OrderableT], y: float) -> Interval[OrderableT]: ... - def __floordiv__(self: Interval[OrderableT], y: float) -> Interval[OrderableT]: ... + @overload + def __add__(self: Interval[OrderableTimesT], y: Timedelta) -> Interval[OrderableTimesT]: ... + @overload + def __add__(self: Interval[int], y: int) -> Interval[int]: ... + @overload + def __add__(self: Interval[int], y: float) -> Interval[float]: ... + @overload + def __add__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... + @overload + def __sub__(self: Interval[OrderableTimesT], y: Timedelta) -> Interval[OrderableTimesT]: ... + @overload + def __sub__(self: Interval[int], y: int) -> Interval[int]: ... + @overload + def __sub__(self: Interval[int], y: float) -> Interval[float]: ... + @overload + def __sub__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... + @overload + def __mul__(self: Interval[int], y: int) -> Interval[int]: ... + @overload + def __mul__(self: Interval[int], y: float) -> Interval[float]: ... + @overload + def __mul__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... + @overload + def __truediv__(self: Interval[int], y: int) -> Interval[int]: ... + @overload + def __truediv__(self: Interval[int], y: float) -> Interval[float]: ... + @overload + def __truediv__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... + @overload + def __floordiv__(self: Interval[int], y: int) -> Interval[int]: ... + @overload + def __floordiv__(self: Interval[int], y: float) -> Interval[float]: ... + @overload + def __floordiv__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... def overlaps(self: Interval[OrderableT], other: Interval[OrderableT]) -> bool: ... diff --git a/partial/pandas/_libs/tslibs/timedeltas.pyi b/partial/pandas/_libs/tslibs/timedeltas.pyi index bafe9eea..d8369f0c 100644 --- a/partial/pandas/_libs/tslibs/timedeltas.pyi +++ b/partial/pandas/_libs/tslibs/timedeltas.pyi @@ -1,84 +1,84 @@ -from __future__ import annotations from datetime import timedelta -from typing import Any, Union, Tuple, Type, Optional, Sequence +from typing import ( + ClassVar, + Type, + TypeVar, + overload, +) import numpy as np -from pandas._typing import Dtype -class _Timedelta(timedelta): - def __hash__(self) -> int: ... - def __richcmp__(self, other, op: int) -> Any: ... - def to_timedelta64(self) -> np.timedelta64: ... - def view(self, dtype: Dtype): ... - @property - def components(self) -> Tuple: ... # Really a namedtuple - @property - def delta(self) -> int: ... - @property - def asm8(self) -> np.timedelta64: ... - @property - def resolution_string(self) -> str: ... - @property - def nanoseconds(self) -> int: ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... - def __bool__(self) -> bool: ... - def isoformat(self) -> str: ... +from pandas._libs.tslibs import ( + NaTType, + Tick, +) +from pandas._typing import npt -class Timedelta(_Timedelta): - def __new__(cls, value: object = ..., unit: Optional[str] = ..., **kwargs) -> Timedelta: ... - def __setstate__(self, state) -> None: ... - def __reduce__(self) -> Tuple[Type[Timedelta], int]: ... - def round(self, freq: str) -> Timedelta: ... - def floor(self, freq: str) -> Timedelta: ... - def ceil(self, freq: str) -> Timedelta: ... - def __inv__(self) -> Timedelta: ... - def __neg__(self) -> Timedelta: ... - def __pos__(self) -> Timedelta: ... - def __abs__(self) -> Timedelta: ... - def __add__(self, other) -> Timedelta: ... - def __radd__(self, other) -> Timedelta: ... - def __sub__(self, other) -> Timedelta: ... - def __rsub_(self, other) -> Timedelta: ... - def __mul__(self, other) -> Timedelta: ... - __rmul__ = __mul__ - def __truediv__(self, other) -> Union[float, Timedelta]: ... - def __rtruediv__(self, other) -> float: ... - def __floordiv__(self, other) -> Union[int, Timedelta]: ... - def __rfloordiv__(self, other) -> int: ... - def __mod__(self, other) -> int: ... - def __rmod__(self, other) -> int: ... - def __divmod__(self, other) -> Tuple[int, int]: ... - def __rdivmod__(self, other) -> Tuple[int, int]: ... +_S = TypeVar("_S", bound=timedelta) - @property - def asm8(self) -> int: ... - @property - def components(self) -> int: ... +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) -> str: ... +def delta_to_nanoseconds(delta: Tick | np.timedelta64 | timedelta | int) -> int: ... + +class Timedelta(timedelta): + min: ClassVar[Timedelta] + max: ClassVar[Timedelta] + resolution: ClassVar[Timedelta] + value: int # np.int64 + + # error: "__new__" must return a class instance (got "Union[Timedelta, NaTType]") + def __new__( # type: ignore[misc] + cls: Type[_S], + value=..., + unit: str = ..., + **kwargs: int | float | np.integer | np.floating, + ) -> _S | NaTType: ... @property def days(self) -> int: ... @property - def delta(self) -> int: ... + def seconds(self) -> int: ... @property def microseconds(self) -> int: ... + def total_seconds(self) -> float: ... + def to_pytimedelta(self) -> timedelta: ... + def to_timedelta64(self) -> np.timedelta64: ... @property - def nanoseconds(self) -> int: ... + 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: ... @property def resolution_string(self) -> str: ... - @property - def seconds(self) -> int: ... - max: Timedelta = ... - min: Timedelta = ... - resolution: 'Timedelta' = ... - def ceil(self, freq, **kwargs) -> Timedelta: ... - def floor(self, freq, **kwargs) -> Timedelta: ... - def isoformat(self) -> str: ... - def round(self, freq) -> Timedelta: ... - def to_numpy(self) -> _np.timedelta64: ... - def to_pytimedelta(self) -> datetime.timedelta: ... - def to_timedelta64(self) -> _np.timedelta64: ... - def total_seconds(self) -> int: ... - - -def delta_to_nanoseconds(delta: _Timedelta) -> int: ... -def ints_to_pytimedelta(arr: Sequence[int], box: bool = ...) -> np.ndarray: ... + def __add__(self, other: timedelta) -> timedelta: ... + def __radd__(self, other: timedelta) -> timedelta: ... + def __sub__(self, other: timedelta) -> timedelta: ... + def __rsub__(self, other: timedelta) -> timedelta: ... + def __neg__(self) -> timedelta: ... + def __pos__(self) -> timedelta: ... + def __abs__(self) -> timedelta: ... + def __mul__(self, other: float) -> timedelta: ... + def __rmul__(self, other: float) -> timedelta: ... + @overload + def __floordiv__(self, other: timedelta) -> int: ... + @overload + def __floordiv__(self, other: int) -> timedelta: ... + @overload + def __truediv__(self, other: timedelta) -> float: ... + @overload + def __truediv__(self, other: float) -> timedelta: ... + def __mod__(self, other: timedelta) -> timedelta: ... + 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: ... + def __hash__(self) -> int: ... diff --git a/partial/pandas/_libs/tslibs/timestamps.pyi b/partial/pandas/_libs/tslibs/timestamps.pyi index 9f9155e6..46bc1bcd 100644 --- a/partial/pandas/_libs/tslibs/timestamps.pyi +++ b/partial/pandas/_libs/tslibs/timestamps.pyi @@ -1,20 +1,110 @@ -from __future__ import annotations -from datetime import datetime, timestamp -import sys -from typing import Any, Optional, Union +from datetime import ( + date as _date, + datetime, + time as _time, + timedelta, + tzinfo as _tzinfo, +) +from time import struct_time +from typing import ( + ClassVar, + TypeVar, + overload, +) -if sys.version_info >= (3, 8): - from typing import Literal -else: - from typing_extensions import Literal +import numpy as np + +from pandas._libs.tslibs import ( + BaseOffset, + NaTType, + Period, + Timedelta, +) + +_DatetimeT = TypeVar("_DatetimeT", bound=datetime) + +def integer_op_not_supported(obj: object) -> TypeError: ... class Timestamp(datetime): - def __init__( + min: ClassVar[Timestamp] + max: ClassVar[Timestamp] + + resolution: ClassVar[Timedelta] + value: int # np.int64 + + # error: "__new__" must return a class instance (got "Union[Timestamp, NaTType]") + def __new__( # type: ignore[misc] + cls: type[_DatetimeT], + ts_input: int | np.integer | float | str | _date | datetime | np.datetime64 = ..., + freq: int | None | str | BaseOffset = ..., + tz: str | _tzinfo | None | int = ..., + unit: str | int | None = ..., + year: int | None = ..., + month: int | None = ..., + day: int | None = ..., + hour: int | None = ..., + minute: int | None = ..., + second: int | None = ..., + microsecond: int | None = ..., + nanosecond: int | None = ..., + tzinfo: _tzinfo | None = ..., + *, + fold: int | None = ..., + ) -> _DatetimeT | NaTType: ... + def _set_freq(self, freq: BaseOffset | None) -> None: ... + @property + def year(self) -> int: ... + @property + def month(self) -> int: ... + @property + def day(self) -> int: ... + @property + def hour(self) -> int: ... + @property + def minute(self) -> int: ... + @property + def second(self) -> int: ... + @property + def microsecond(self) -> int: ... + @property + def tzinfo(self) -> _tzinfo | None: ... + @property + def tz(self) -> _tzinfo | None: ... + @property + def fold(self) -> int: ... + @classmethod + def fromtimestamp(cls: type[_DatetimeT], t: float, tz: _tzinfo | None = ...) -> _DatetimeT: ... + @classmethod + def utcfromtimestamp(cls: type[_DatetimeT], t: float) -> _DatetimeT: ... + @classmethod + def today(cls: type[_DatetimeT], tz: _tzinfo | str | None = ...) -> _DatetimeT: ... + @classmethod + def fromordinal( + cls: type[_DatetimeT], + ordinal: int, + freq: str | BaseOffset | None = ..., + tz: _tzinfo | str | None = ..., + ) -> _DatetimeT: ... + @classmethod + def now(cls: type[_DatetimeT], tz: _tzinfo | str | None = ...) -> _DatetimeT: ... + @classmethod + def utcnow(cls: type[_DatetimeT]) -> _DatetimeT: ... + # error: Signature of "combine" incompatible with supertype "datetime" + @classmethod + def combine(cls, date: _date, time: _time) -> datetime: ... # type: ignore[override] + @classmethod + def fromisoformat(cls: type[_DatetimeT], date_string: str) -> _DatetimeT: ... + def strftime(self, format: str) -> str: ... + def __format__(self, fmt: str) -> str: ... + def toordinal(self) -> int: ... + def timetuple(self) -> struct_time: ... + def timestamp(self) -> float: ... + def utctimetuple(self) -> struct_time: ... + def date(self) -> _date: ... + def time(self) -> _time: ... + def timetz(self) -> _time: ... + def replace( self, - ts_input: Any, - freq: Any = ..., - tz: Any = ..., - unit: str = ..., year: int = ..., month: int = ..., day: int = ..., @@ -22,144 +112,82 @@ class Timestamp(datetime): minute: int = ..., second: int = ..., microsecond: int = ..., - nanosecond: int = ..., - tzinfo: Any = ..., - ): ... - - # Class methods - # Methods - -class Tick(SingleConstructorOffset): - __array_priority__: int - def __init__(self, n: int = ..., normalize: bool = ...) -> None: ... - @property - def delta(self) -> int: ... - @property - def nanos(self) -> int: ... - def is_on_offset(self, dt: datetime) -> bool: ... - def is_anchored(self) -> bool: ... - def __eq__(self, other) -> bool: ... - def __ne__(self, other) -> bool: ... - def __le__(self, other) -> bool: ... - def __lt__(self, other) -> bool: ... - def __ge__(self, other) -> bool: ... - def __gt__(self, other) -> bool: ... - def __mul__(self, other) -> Tick: ... - def __truediv__(self, other) -> Tick: ... - def __add__(self, other) -> Tick: ... - def apply(self, other) -> Any: ... - def __setstate__(self, state: Mapping) -> None: ... - -class Timestamp(timestamp): - @staticmethod - def combine(date, time) -> Timestamp: ... - @staticmethod - def fromordinal(ordinal: int, freq: Any = ..., tz: Any = ...) -> Timestamp: ... - @staticmethod - def fromtimestamp(ts) -> Timestamp: ... - @staticmethod - def now(tz = ...) -> Timestamp: ... - @staticmethod - def today(tz = ...) -> Timestamp: ... - @staticmethod - def utcfromtimestamp(ts) -> Timestamp: ... - @staticmethod - def utcnow() -> Timestamp: ... - - def normalize(self) -> Timestamp: ... - def replace( - self, - year: Optional[int] = ..., - month: Optional[int] = ..., - day: Optional[int] = ..., - hour: Optional[int] = ..., - minute: Optional[int] = ..., - second: Optional[int] = ..., - microsecond: Optional[int] = ..., - nanosecond: Optional[int] = ..., - tzinfo: Any = ..., + tzinfo: _tzinfo | None = ..., fold: int = ..., - ) -> Timestamp: ... - - def astimezone(self, tz: Any) -> Timestamp: ... + ) -> datetime: ... + def astimezone(self: _DatetimeT, tz: _tzinfo | None = ...) -> _DatetimeT: ... def ctime(self) -> str: ... - def date(self) -> Any: ... - def dst(self) -> Timestamp: ... - def fromisoformat(self) -> Any: ... - def isocalendar(self) -> Tuple[int, int, int]: ... - def isoweekday(self) -> int: ... - def monthname(self, local: Any) -> str: ... - def day_name(self, local: Optional[str] = ...) -> str: ... - - def ceil(self, freq: str, ambiguous: Union[bool, str, Literal['raise', 'NaT']] = ..., - nonexistent: Union[Timedelta, str, Literal['raise', 'shift_forward', 'shift_backward', 'NaT']] = ...) \ - -> Timestamp: ... - def floor(self, freq: str, ambiguous: Union[bool, str, Literal['raise', 'NaT']] = ..., - nonexistent: Union[Timedelta, str, Literal['raise', 'shift_forward', 'shift_backward', 'NaT']] = ...) \ - -> Timestamp: ... - def round(self, freq: str, ambiguous: Union[bool, str, Literal['raise', 'NaT']] = ..., - nonexistent: Union[Timedelta, str, Literal['raise', 'shift_forward', 'shift_backward', 'NaT']] = ...) \ - -> Timestamp: ... - def to_julian_date(self, *args, **kwargs) -> str: ... - def strftime(self, fmt: str) -> str: ... - def time(self) -> Timestamp: ... - def timestamp(self) -> float: ... - def timetuple(self) -> Tuple: ... - def timetz(self) -> Timestamp: ... - def to_datetime64(self) -> _np.datetime64: ... - def to_numpy(self) -> _np.datetime64: ... - def to_period(self, freq: Optional[str] = ...) -> Any: ... - def to_pydatetime(self, warn: bool = ...) -> datetime: ... - def toordinal(self) -> Any: ... - def tz_convert(self, tz: Any) -> Timestamp: ... - def tz_localize(self, tz: Any, ambiguous: Any = ..., nonexistent: Any = ...) -> Timestamp: ... - def tzname(self) -> str: ... - def utcoffset(self) -> Any: ... - def utctimetuple(self) -> Any: ... - def weekday(self) -> int: ... - + def isoformat(self, sep: str = ..., timespec: str = ...) -> str: ... @classmethod - def utcfromtimestamp(cls, ts: Timestamp) -> Timestamp: ... - @classmethod - def utcnow(cls) -> Timestamp: ... - def __init__(self, *args, **kwargs): ... - def __new__(cls, *args, **kwargs) -> Timestamp: ... - - # Still need to do parameter types - @property - def asm8(self) -> int: ... - @property - def tz(self) -> Any: ... - @property - def dayofweek(self) -> int: ... - @property - def dayofyear(self) -> int: ... + def strptime(cls, date_string: str, format: str) -> datetime: ... + def utcoffset(self) -> timedelta | None: ... + def tzname(self) -> str | None: ... + def dst(self) -> timedelta | None: ... + def __le__(self, other: datetime) -> bool: ... # type: ignore + def __lt__(self, other: datetime) -> bool: ... # type: ignore + def __ge__(self, other: datetime) -> bool: ... # type: ignore + def __gt__(self, other: datetime) -> bool: ... # type: ignore + # error: Signature of "__add__" incompatible with supertype "date"/"datetime" + @overload # type: ignore[override] + def __add__(self, other: np.ndarray) -> np.ndarray: ... + @overload + # TODO: other can also be Tick (but it cannot be resolved) + def __add__(self: _DatetimeT, other: timedelta | np.timedelta64) -> _DatetimeT: ... + def __radd__(self: _DatetimeT, other: timedelta) -> _DatetimeT: ... + @overload # type: ignore + def __sub__(self, other: Timestamp) -> Timedelta: ... + @overload # type: ignore + def __sub__(self, other: datetime) -> timedelta: ... + @overload + # TODO: other can also be Tick (but it cannot be resolved) + def __sub__(self, other: timedelta | np.timedelta64) -> datetime: ... + def __hash__(self) -> int: ... + def weekday(self) -> int: ... + def isoweekday(self) -> int: ... + def isocalendar(self) -> tuple[int, int, int]: ... @property - def daysinmonth(self) -> int: ... + def is_leap_year(self) -> bool: ... @property - def days_in_month(self) -> int: ... + def is_month_start(self) -> bool: ... @property - def freqstr(self) -> str: ... + def is_quarter_start(self) -> bool: ... @property - def is_leap_year(self) -> bool: ... + def is_year_start(self) -> bool: ... @property def is_month_end(self) -> bool: ... @property - def is_month_start(self) -> bool: ... - @property def is_quarter_end(self) -> bool: ... @property - def is_quarter_start(self) -> bool: ... - @property def is_year_end(self) -> bool: ... - @property - def is_year_start(self) -> bool: ... + def to_pydatetime(self, warn: bool = ...) -> datetime: ... + def to_datetime64(self) -> np.datetime64: ... + def to_period(self, freq: BaseOffset | str | None = ...) -> Period: ... + def to_julian_date(self) -> np.float64: ... + @property + def asm8(self) -> np.datetime64: ... + def tz_convert(self: _DatetimeT, tz: _tzinfo | str | None) -> _DatetimeT: ... + # TODO: could return NaT? + def tz_localize( + self: _DatetimeT, + tz: _tzinfo | str | None, + ambiguous: str = ..., + nonexistent: str = ..., + ) -> _DatetimeT: ... + def normalize(self: _DatetimeT) -> _DatetimeT: ... + # TODO: round/floor/ceil could return NaT? + def round(self: _DatetimeT, freq: str, ambiguous: bool | str = ..., nonexistent: str = ...) -> _DatetimeT: ... + def floor(self: _DatetimeT, freq: str, ambiguous: bool | str = ..., nonexistent: str = ...) -> _DatetimeT: ... + def ceil(self: _DatetimeT, freq: str, ambiguous: bool | str = ..., nonexistent: str = ...) -> _DatetimeT: ... + def day_name(self, locale: str | None = ...) -> str: ... + def month_name(self, locale: str | None = ...) -> str: ... + @property + def day_of_week(self) -> int: ... + @property + def day_of_month(self) -> int: ... + @property + def day_of_year(self) -> int: ... @property def quarter(self) -> int: ... @property def week(self) -> int: ... - @property - def weekofyear(self) -> int: ... - -from .timedeltas import Timedelta -from .period import Period + def to_numpy(self, dtype: np.dtype | None = ..., copy: bool = ...) -> np.datetime64: ... diff --git a/tests/pandas/test_interval.py b/tests/pandas/test_interval.py index c637d065..30becfbe 100644 --- a/tests/pandas/test_interval.py +++ b/tests/pandas/test_interval.py @@ -1,7 +1,12 @@ # flake8: noqa: F841 - +from typing import TYPE_CHECKING import pandas as pd +if not TYPE_CHECKING: + + def reveal_type(arg, **kwargs): + pass + def test_interval_init() -> None: i1: pd.Interval = pd.Interval(1, 2, closed="both") @@ -27,12 +32,43 @@ def test_max_intervals() -> None: def test_interval_length() -> None: - i1 = pd.Interval(pd.Timestamp("2000-01-01"), pd.Timestamp("2000-01-02"), closed="both") + i1 = pd.Interval(pd.Timestamp("2000-01-01"), pd.Timestamp("2000-01-03"), closed="both") reveal_type(i1.length, expected_string="Timedelta") + reveal_type(i1.left, expected_string="Timestamp") + reveal_type(i1.right, expected_string="Timestamp") i1.length.total_seconds() + pd.Timestamp("2001-01-02") in i1 + i1 + pd.Timedelta(seconds=20) + + reveal_type(i1 + pd.Timedelta(seconds=20), expected_string="Interval[Timestamp]") + if TYPE_CHECKING: + 20 in i1 # type: ignore + i1 + pd.Timestamp("2000-03-03") # type: ignore + i1 * 3 # type: ignore + i1 * pd.Timedelta(seconds=20) # type: ignore i2 = pd.Interval(10, 20) reveal_type(i2.length, expected_type=int) + reveal_type(i2.left, expected_type=int) + reveal_type(i2.right, expected_type=int) + 15 in i2 + reveal_type(i2 + 3, expected_string="Interval[int]") + reveal_type(i2 + 3.2, expected_string="Interval[float]") + reveal_type(i2 * 4, expected_string="Interval[int]") + reveal_type(i2 * 4.2, expected_string="Interval[float]") + + if TYPE_CHECKING: + pd.Timestamp("2001-01-02") in i2 # type: ignore + i2 + pd.Timedelta(seconds=20) # type: ignore + reveal_type(pd.Timestamp("2001-02-02")) i3 = pd.Interval(13.2, 19.5) reveal_type(i3.length, expected_type=float) + reveal_type(i3.left, expected_type=float) + reveal_type(i3.right, expected_type=float) + 15.4 in i3 + reveal_type(i3 + 3, expected_string="Interval[float]") + reveal_type(i3 * 3, expected_string="Interval[float]") + if TYPE_CHECKING: + pd.Timestamp("2001-01-02") in i3 # type: ignore + i3 + pd.Timedelta(seconds=20) # type: ignore From a57ccbf1101dcfa2e4f09b8824c424e2a97e71d1 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Mon, 21 Feb 2022 08:09:54 -0500 Subject: [PATCH 4/8] remove new use init --- partial/pandas/_libs/interval.pyi | 31 +++---------------------------- 1 file changed, 3 insertions(+), 28 deletions(-) diff --git a/partial/pandas/_libs/interval.pyi b/partial/pandas/_libs/interval.pyi index a8a969bb..fed91029 100644 --- a/partial/pandas/_libs/interval.pyi +++ b/partial/pandas/_libs/interval.pyi @@ -43,34 +43,9 @@ class Interval(IntervalMixin, Generic[OrderableT]): def right(self: Interval[OrderableT]) -> OrderableT: ... @property def closed(self) -> str: ... - @overload - def __new__( - cls, - left: Timestamp, - right: Timestamp, - closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., - ) -> Interval[Timestamp]: ... - @overload - def __new__( - cls, - left: Timedelta, - right: Timedelta, - closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., - ) -> Interval[Timedelta]: ... - @overload - def __new__( - cls, - left: int, - right: int, - closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., - ) -> Interval[int]: ... - @overload - def __new__( - cls, - left: float, - right: float, - closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., - ) -> Interval[float]: ... + def __init__( + self, left: OrderableT, right: OrderableT, closed: Union[str, Literal["left", "right", "both", "neither"]] = ... + ): ... length: _LengthDescriptor def __hash__(self) -> int: ... @overload From 655928faee61dc1ad190ed738a8f07de55383e30 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Mon, 21 Feb 2022 12:24:27 -0500 Subject: [PATCH 5/8] further refinement - now works with pandas mypy check --- partial/pandas/_libs/interval.pyi | 44 +++++++++++++++++++++++++------ 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/partial/pandas/_libs/interval.pyi b/partial/pandas/_libs/interval.pyi index fed91029..9a7cd5d1 100644 --- a/partial/pandas/_libs/interval.pyi +++ b/partial/pandas/_libs/interval.pyi @@ -1,13 +1,18 @@ from __future__ import annotations import sys from typing import Generic, overload, Union, Protocol, TypeVar, Any -from _typing import Timedelta, Timestamp + +import numpy as np + +from pandas._typing import Timedelta, Timestamp if sys.version_info >= (3, 8): from typing import Literal else: from typing_extensions import Literal +VALID_CLOSED: frozenset + OrderableScalarT = TypeVar("OrderableScalarT", int, float) OrderableTimesT = TypeVar("OrderableTimesT", Timestamp, Timedelta) OrderableT = TypeVar("OrderableT", int, float, Timestamp, Timedelta) @@ -35,6 +40,7 @@ class IntervalMixin: def mid(self: IntervalMixinProtocol) -> float: ... @property def is_empty(self: IntervalMixinProtocol) -> bool: ... + def _check_closed_matches(self, other: IntervalMixin, name: str = ...): ... class Interval(IntervalMixin, Generic[OrderableT]): @property @@ -44,14 +50,15 @@ class Interval(IntervalMixin, Generic[OrderableT]): @property def closed(self) -> str: ... def __init__( - self, left: OrderableT, right: OrderableT, closed: Union[str, Literal["left", "right", "both", "neither"]] = ... + self, + left: OrderableT, + right: OrderableT, + closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., ): ... length: _LengthDescriptor def __hash__(self) -> int: ... @overload - def __contains__(self: Interval[Timestamp], key: Timestamp) -> bool: ... - @overload - def __contains__(self: Interval[Timedelta], key: Timedelta) -> bool: ... + def __contains__(self: Interval[OrderableTimesT], OrderableTimesT) -> bool: ... @overload def __contains__(self: Interval[int], key: Union[int, float]) -> bool: ... @overload @@ -59,7 +66,9 @@ class Interval(IntervalMixin, Generic[OrderableT]): def __repr__(self) -> str: ... def __str__(self) -> str: ... @overload - def __add__(self: Interval[OrderableTimesT], y: Timedelta) -> Interval[OrderableTimesT]: ... + def __add__( + self: Interval[OrderableTimesT], y: Timedelta + ) -> Interval[OrderableTimesT]: ... @overload def __add__(self: Interval[int], y: int) -> Interval[int]: ... @overload @@ -67,7 +76,9 @@ class Interval(IntervalMixin, Generic[OrderableT]): @overload def __add__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... @overload - def __sub__(self: Interval[OrderableTimesT], y: Timedelta) -> Interval[OrderableTimesT]: ... + def __sub__( + self: Interval[OrderableTimesT], y: Timedelta + ) -> Interval[OrderableTimesT]: ... @overload def __sub__(self: Interval[int], y: int) -> Interval[int]: ... @overload @@ -91,5 +102,22 @@ class Interval(IntervalMixin, Generic[OrderableT]): @overload def __floordiv__(self: Interval[int], y: float) -> Interval[float]: ... @overload - def __floordiv__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... + def __floordiv__( + self: Interval[float], y: Union[int, float] + ) -> Interval[float]: ... def overlaps(self: Interval[OrderableT], other: Interval[OrderableT]) -> bool: ... + +def intervals_to_interval_bounds(intervals: np.ndarray, validate_closed: int = ...): ... + +class IntervalTree(IntervalMixin): + def __init__( + self, + left: np.ndarray, + right: np.ndarray, + closed: Literal["left", "right", "both", "neither"] = ..., + ): ... + def get_indexer(self, target) -> np.ndarray: ... + def get_indexer_non_unique(self, target) -> np.ndarray: ... + _na_count: int + @property + def is_overlapping(self) -> bool: ... From b2ee28263cee40cbd13bc2a71ec5c021a5a809c8 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Mon, 21 Feb 2022 14:59:51 -0500 Subject: [PATCH 6/8] feedback on pandas PR --- partial/pandas/_libs/interval.pyi | 54 +++++++++++++++++++------------ tests/pandas/test_interval.py | 5 +++ 2 files changed, 38 insertions(+), 21 deletions(-) diff --git a/partial/pandas/_libs/interval.pyi b/partial/pandas/_libs/interval.pyi index 9a7cd5d1..516b91f5 100644 --- a/partial/pandas/_libs/interval.pyi +++ b/partial/pandas/_libs/interval.pyi @@ -1,10 +1,21 @@ from __future__ import annotations + import sys -from typing import Generic, overload, Union, Protocol, TypeVar, Any +from typing import ( + Any, + Generic, + Protocol, + TypeVar, + Union, + overload, +) import numpy as np -from pandas._typing import Timedelta, Timestamp +from pandas._typing import ( + Timedelta, + Timestamp, +) if sys.version_info >= (3, 8): from typing import Literal @@ -27,19 +38,27 @@ class _LengthDescriptor: @overload def __get__(self, instance: Interval[OrderableTimesT], owner: Any) -> Timedelta: ... -class IntervalMixin: - @property - def closed_left(self: IntervalMixinProtocol) -> bool: ... +class _MidDescriptor: + @overload + def __get__(self, instance: Interval[OrderableScalarT], owner: Any) -> float: ... + @overload + def __get__(self, instance: Interval[Timedelta], owner: Any) -> Timedelta: ... + @overload + def __get__(self, instance: Interval[Timestamp], owner: Any) -> Timestamp: ... + +class IntervalMixin(IntervalMixinProtocol): @property - def closed_right(self: IntervalMixinProtocol) -> bool: ... + def closed_left(self) -> bool: ... @property - def open_left(self: IntervalMixinProtocol) -> bool: ... + def closed_right(self) -> bool: ... @property - def open_right(self: IntervalMixinProtocol) -> bool: ... + def open_left(self) -> bool: ... @property - def mid(self: IntervalMixinProtocol) -> float: ... + def open_right(self) -> bool: ... + mid: _MidDescriptor + length: _LengthDescriptor @property - def is_empty(self: IntervalMixinProtocol) -> bool: ... + def is_empty(self) -> bool: ... def _check_closed_matches(self, other: IntervalMixin, name: str = ...): ... class Interval(IntervalMixin, Generic[OrderableT]): @@ -53,9 +72,8 @@ class Interval(IntervalMixin, Generic[OrderableT]): self, left: OrderableT, right: OrderableT, - closed: Union[str, Literal["left", "right", "both", "neither"]] = ..., + closed: Literal["left", "right", "both", "neither"] = ..., ): ... - length: _LengthDescriptor def __hash__(self) -> int: ... @overload def __contains__(self: Interval[OrderableTimesT], OrderableTimesT) -> bool: ... @@ -66,9 +84,7 @@ class Interval(IntervalMixin, Generic[OrderableT]): def __repr__(self) -> str: ... def __str__(self) -> str: ... @overload - def __add__( - self: Interval[OrderableTimesT], y: Timedelta - ) -> Interval[OrderableTimesT]: ... + def __add__(self: Interval[OrderableTimesT], y: Timedelta) -> Interval[OrderableTimesT]: ... @overload def __add__(self: Interval[int], y: int) -> Interval[int]: ... @overload @@ -76,9 +92,7 @@ class Interval(IntervalMixin, Generic[OrderableT]): @overload def __add__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... @overload - def __sub__( - self: Interval[OrderableTimesT], y: Timedelta - ) -> Interval[OrderableTimesT]: ... + def __sub__(self: Interval[OrderableTimesT], y: Timedelta) -> Interval[OrderableTimesT]: ... @overload def __sub__(self: Interval[int], y: int) -> Interval[int]: ... @overload @@ -102,9 +116,7 @@ class Interval(IntervalMixin, Generic[OrderableT]): @overload def __floordiv__(self: Interval[int], y: float) -> Interval[float]: ... @overload - def __floordiv__( - self: Interval[float], y: Union[int, float] - ) -> Interval[float]: ... + def __floordiv__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... def overlaps(self: Interval[OrderableT], other: Interval[OrderableT]) -> bool: ... def intervals_to_interval_bounds(intervals: np.ndarray, validate_closed: int = ...): ... diff --git a/tests/pandas/test_interval.py b/tests/pandas/test_interval.py index 30becfbe..97b25e6e 100644 --- a/tests/pandas/test_interval.py +++ b/tests/pandas/test_interval.py @@ -36,6 +36,7 @@ def test_interval_length() -> None: reveal_type(i1.length, expected_string="Timedelta") reveal_type(i1.left, expected_string="Timestamp") reveal_type(i1.right, expected_string="Timestamp") + reveal_type(i1.mid, expected_string="Timestamp") i1.length.total_seconds() pd.Timestamp("2001-01-02") in i1 i1 + pd.Timedelta(seconds=20) @@ -51,6 +52,8 @@ def test_interval_length() -> None: reveal_type(i2.length, expected_type=int) reveal_type(i2.left, expected_type=int) reveal_type(i2.right, expected_type=int) + reveal_type(i2.mid, expected_type=float) + 15 in i2 reveal_type(i2 + 3, expected_string="Interval[int]") reveal_type(i2 + 3.2, expected_string="Interval[float]") @@ -66,6 +69,8 @@ def test_interval_length() -> None: reveal_type(i3.length, expected_type=float) reveal_type(i3.left, expected_type=float) reveal_type(i3.right, expected_type=float) + reveal_type(i3.mid, expected_type=float) + 15.4 in i3 reveal_type(i3 + 3, expected_string="Interval[float]") reveal_type(i3 * 3, expected_string="Interval[float]") From de048ae05753f0d738e258f3894f4f3fd2e1c246 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Mon, 21 Feb 2022 21:51:03 -0500 Subject: [PATCH 7/8] feedback from pandas core team review --- partial/pandas/_libs/interval.pyi | 118 ++++++++++++++++++++---------- 1 file changed, 81 insertions(+), 37 deletions(-) diff --git a/partial/pandas/_libs/interval.pyi b/partial/pandas/_libs/interval.pyi index 516b91f5..2ad2cf79 100644 --- a/partial/pandas/_libs/interval.pyi +++ b/partial/pandas/_libs/interval.pyi @@ -1,34 +1,28 @@ from __future__ import annotations -import sys from typing import ( Any, Generic, - Protocol, + Literal, + Tuple, TypeVar, Union, overload, ) import numpy as np +import numpy.typing as npt from pandas._typing import ( Timedelta, Timestamp, ) -if sys.version_info >= (3, 8): - from typing import Literal -else: - from typing_extensions import Literal +VALID_CLOSED: frozenset[str] -VALID_CLOSED: frozenset - -OrderableScalarT = TypeVar("OrderableScalarT", int, float) -OrderableTimesT = TypeVar("OrderableTimesT", Timestamp, Timedelta) -OrderableT = TypeVar("OrderableT", int, float, Timestamp, Timedelta) - -class IntervalMixinProtocol(Protocol): ... +_OrderableScalarT = TypeVar("_OrderableScalarT", int, float) +_OrderableTimesT = TypeVar("_OrderableTimesT", Timestamp, Timedelta) +_OrderableT = TypeVar("_OrderableT", int, float, Timestamp, Timedelta) class _LengthDescriptor: @overload @@ -36,17 +30,27 @@ class _LengthDescriptor: @overload def __get__(self, instance: Interval[int], owner: Any) -> int: ... @overload - def __get__(self, instance: Interval[OrderableTimesT], owner: Any) -> Timedelta: ... + def __get__( + self, instance: Interval[_OrderableTimesT], owner: Any + ) -> Timedelta: ... + @overload + def __get__( + self, instance: IntervalTree[_OrderableT], owner: Any + ) -> _OrderableT: ... class _MidDescriptor: @overload - def __get__(self, instance: Interval[OrderableScalarT], owner: Any) -> float: ... + def __get__(self, instance: Interval[_OrderableScalarT], owner: Any) -> float: ... @overload - def __get__(self, instance: Interval[Timedelta], owner: Any) -> Timedelta: ... + def __get__( + self, instance: Interval[_OrderableTimesT], owner: Any + ) -> _OrderableTimesT: ... @overload - def __get__(self, instance: Interval[Timestamp], owner: Any) -> Timestamp: ... + def __get__( + self, instance: IntervalTree[_OrderableT], owner: Any + ) -> _OrderableT: ... -class IntervalMixin(IntervalMixinProtocol): +class IntervalMixin: @property def closed_left(self) -> bool: ... @property @@ -59,32 +63,34 @@ class IntervalMixin(IntervalMixinProtocol): length: _LengthDescriptor @property def is_empty(self) -> bool: ... - def _check_closed_matches(self, other: IntervalMixin, name: str = ...): ... + def _check_closed_matches(self, other: IntervalMixin, name: str = ...) -> None: ... -class Interval(IntervalMixin, Generic[OrderableT]): +class Interval(IntervalMixin, Generic[_OrderableT]): @property - def left(self: Interval[OrderableT]) -> OrderableT: ... + def left(self: Interval[_OrderableT]) -> _OrderableT: ... @property - def right(self: Interval[OrderableT]) -> OrderableT: ... + def right(self: Interval[_OrderableT]) -> _OrderableT: ... @property - def closed(self) -> str: ... + def closed(self) -> Literal["left", "right", "both", "neither"]: ... def __init__( self, - left: OrderableT, - right: OrderableT, + left: _OrderableT, + right: _OrderableT, closed: Literal["left", "right", "both", "neither"] = ..., ): ... def __hash__(self) -> int: ... @overload - def __contains__(self: Interval[OrderableTimesT], OrderableTimesT) -> bool: ... - @overload - def __contains__(self: Interval[int], key: Union[int, float]) -> bool: ... + def __contains__(self: Interval[_OrderableTimesT], _OrderableTimesT) -> bool: ... @overload - def __contains__(self: Interval[float], key: Union[int, float]) -> bool: ... + def __contains__( + self: Interval[_OrderableScalarT], key: Union[int, float] + ) -> bool: ... def __repr__(self) -> str: ... def __str__(self) -> str: ... @overload - def __add__(self: Interval[OrderableTimesT], y: Timedelta) -> Interval[OrderableTimesT]: ... + def __add__( + self: Interval[_OrderableTimesT], y: Timedelta + ) -> Interval[_OrderableTimesT]: ... @overload def __add__(self: Interval[int], y: int) -> Interval[int]: ... @overload @@ -92,7 +98,19 @@ class Interval(IntervalMixin, Generic[OrderableT]): @overload def __add__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... @overload - def __sub__(self: Interval[OrderableTimesT], y: Timedelta) -> Interval[OrderableTimesT]: ... + def __radd__( + self: Interval[_OrderableTimesT], y: Timedelta + ) -> Interval[_OrderableTimesT]: ... + @overload + def __radd__(self: Interval[int], y: int) -> Interval[int]: ... + @overload + def __radd__(self: Interval[int], y: float) -> Interval[float]: ... + @overload + def __radd__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... + @overload + def __sub__( + self: Interval[_OrderableTimesT], y: Timedelta + ) -> Interval[_OrderableTimesT]: ... @overload def __sub__(self: Interval[int], y: int) -> Interval[int]: ... @overload @@ -100,12 +118,28 @@ class Interval(IntervalMixin, Generic[OrderableT]): @overload def __sub__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... @overload + def __rsub__( + self: Interval[_OrderableTimesT], y: Timedelta + ) -> Interval[_OrderableTimesT]: ... + @overload + def __rsub__(self: Interval[int], y: int) -> Interval[int]: ... + @overload + def __rsub__(self: Interval[int], y: float) -> Interval[float]: ... + @overload + def __rsub__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... + @overload def __mul__(self: Interval[int], y: int) -> Interval[int]: ... @overload def __mul__(self: Interval[int], y: float) -> Interval[float]: ... @overload def __mul__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... @overload + def __rmul__(self: Interval[int], y: int) -> Interval[int]: ... + @overload + def __rmul__(self: Interval[int], y: float) -> Interval[float]: ... + @overload + def __rmul__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... + @overload def __truediv__(self: Interval[int], y: int) -> Interval[int]: ... @overload def __truediv__(self: Interval[int], y: float) -> Interval[float]: ... @@ -116,20 +150,30 @@ class Interval(IntervalMixin, Generic[OrderableT]): @overload def __floordiv__(self: Interval[int], y: float) -> Interval[float]: ... @overload - def __floordiv__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ... - def overlaps(self: Interval[OrderableT], other: Interval[OrderableT]) -> bool: ... + def __floordiv__( + self: Interval[float], y: Union[int, float] + ) -> Interval[float]: ... + def overlaps(self: Interval[_OrderableT], other: Interval[_OrderableT]) -> bool: ... -def intervals_to_interval_bounds(intervals: np.ndarray, validate_closed: int = ...): ... +def intervals_to_interval_bounds( + intervals: np.ndarray, validate_closed: bool = ... +) -> Tuple[np.ndarray, np.ndarray, str]: ... -class IntervalTree(IntervalMixin): +class IntervalTree(IntervalMixin, Generic[_OrderableT]): def __init__( self, left: np.ndarray, right: np.ndarray, closed: Literal["left", "right", "both", "neither"] = ..., + leaf_size: int = ..., ): ... - def get_indexer(self, target) -> np.ndarray: ... - def get_indexer_non_unique(self, target) -> np.ndarray: ... + def get_indexer(self, target) -> npt.NDArray[np.intp]: ... + def get_indexer_non_unique( + self, target + ) -> Tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ... _na_count: int @property def is_overlapping(self) -> bool: ... + @property + def is_monotonic_increasing(self) -> bool: ... + def clear_mapping(self) -> None: ... From 58fb32af159a93a3f3a109debfb2145564a7c638 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Tue, 22 Feb 2022 12:14:33 -0500 Subject: [PATCH 8/8] remove generic from IntervalTree --- partial/pandas/_libs/interval.pyi | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/partial/pandas/_libs/interval.pyi b/partial/pandas/_libs/interval.pyi index 2ad2cf79..77eeda14 100644 --- a/partial/pandas/_libs/interval.pyi +++ b/partial/pandas/_libs/interval.pyi @@ -34,9 +34,7 @@ class _LengthDescriptor: self, instance: Interval[_OrderableTimesT], owner: Any ) -> Timedelta: ... @overload - def __get__( - self, instance: IntervalTree[_OrderableT], owner: Any - ) -> _OrderableT: ... + def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ... class _MidDescriptor: @overload @@ -46,9 +44,7 @@ class _MidDescriptor: self, instance: Interval[_OrderableTimesT], owner: Any ) -> _OrderableTimesT: ... @overload - def __get__( - self, instance: IntervalTree[_OrderableT], owner: Any - ) -> _OrderableT: ... + def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ... class IntervalMixin: @property @@ -159,7 +155,7 @@ def intervals_to_interval_bounds( intervals: np.ndarray, validate_closed: bool = ... ) -> Tuple[np.ndarray, np.ndarray, str]: ... -class IntervalTree(IntervalMixin, Generic[_OrderableT]): +class IntervalTree(IntervalMixin): def __init__( self, left: np.ndarray,