pandas: get types of Interval correct

partial/pandas/_libs/interval.pyi

@@ -1,71 +1,175 @@
 from __future__ import annotations
-import sys
-from typing import Generic, overload, Union, Protocol, TypeVar
-from _typing import Timedelta, Timestamp
-import datetime
 
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal
+from typing import (
+    Any,
+    Generic,
+    Literal,
+    Tuple,
+    TypeVar,
+    Union,
+    overload,
+)
 
-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)
+import numpy as np
+import numpy.typing as npt
 
-class IntervalMixinProtocol(Protocol): ...
+from pandas._typing import (
+    Timedelta,
+    Timestamp,
+)
+
+VALID_CLOSED: frozenset[str]
+
+_OrderableScalarT = TypeVar("_OrderableScalarT", int, float)
+_OrderableTimesT = TypeVar("_OrderableTimesT", Timestamp, Timedelta)
+_OrderableT = TypeVar("_OrderableT", int, float, Timestamp, Timedelta)
+
+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[_OrderableTimesT], owner: Any
+    ) -> Timedelta: ...
+    @overload
+    def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ...
+
+class _MidDescriptor:
+    @overload
+    def __get__(self, instance: Interval[_OrderableScalarT], owner: Any) -> float: ...
+    @overload
+    def __get__(
+        self, instance: Interval[_OrderableTimesT], owner: Any
+    ) -> _OrderableTimesT: ...
+    @overload
+    def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ...
 
 class IntervalMixin:
     @property
-    def closed_left(self: IntervalMixinProtocol) -> bool: ...
+    def closed_left(self) -> bool: ...
     @property
-    def closed_right(self: IntervalMixinProtocol) -> bool: ...
+    def closed_right(self) -> bool: ...
     @property
-    def open_left(self: IntervalMixinProtocol) -> bool: ...
+    def open_left(self) -> bool: ...
     @property
-    def open_right(self: IntervalMixinProtocol) -> bool: ...
+    def open_right(self) -> bool: ...
+    mid: _MidDescriptor
+    length: _LengthDescriptor
     @property
-    def mid(self: IntervalMixinProtocol) -> float: ...
-    @property
-    def is_empty(self: IntervalMixinProtocol) -> bool: ...
+    def is_empty(self) -> bool: ...
+    def _check_closed_matches(self, other: IntervalMixin, name: str = ...) -> None: ...
 
-class Interval(IntervalMixin, Generic[Orderable]):
-    @overload
-    def left(self: Interval[OrderableScalar]) -> OrderableScalar: ...
-    @overload
-    def left(self: Interval[OrderableTimes]) -> OrderableTimes: ...
+class Interval(IntervalMixin, Generic[_OrderableT]):
     @property
-    @overload
-    def left(self: Interval[Orderable]) -> Orderable: ...
-    @overload
-    def right(self: Interval[OrderableScalar]) -> OrderableScalar: ...
-    @overload
-    def right(self: Interval[OrderableTimes]) -> OrderableTimes: ...
+    def left(self: Interval[_OrderableT]) -> _OrderableT: ...
     @property
-    @overload
-    def right(self: Interval[Orderable]) -> Orderable: ...
+    def right(self: Interval[_OrderableT]) -> _OrderableT: ...
     @property
-    def closed(self) -> str: ...
+    def closed(self) -> Literal["left", "right", "both", "neither"]: ...
     def __init__(
         self,
-        left: Orderable,
-        right: Orderable,
-        closed: Union[str, Literal["left", "right", "both", "neither"]] = ...,
-    ) -> None: ...
-    @overload
-    def length(self: Interval[OrderableScalar]) -> float: ...
+        left: _OrderableT,
+        right: _OrderableT,
+        closed: Literal["left", "right", "both", "neither"] = ...,
+    ): ...
+    def __hash__(self) -> int: ...
     @overload
-    def length(self: Interval[OrderableTimes]) -> Timedelta: ...
-    @property
+    def __contains__(self: Interval[_OrderableTimesT], _OrderableTimesT) -> bool: ...
     @overload
-    def length(self: Interval[Orderable]) -> Orderable: ...
-    def __hash__(self) -> int: ...
-    def __contains__(self, key: Orderable) -> bool: ...
+    def __contains__(
+        self: Interval[_OrderableScalarT], key: Union[int, float]
+    ) -> 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: ...
+    @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 __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
+    def __sub__(self: Interval[int], y: float) -> Interval[float]: ...
+    @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]: ...
+    @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: ...
+
+def intervals_to_interval_bounds(
+    intervals: np.ndarray, validate_closed: bool = ...
+) -> Tuple[np.ndarray, np.ndarray, str]: ...
+
+class IntervalTree(IntervalMixin):
+    def __init__(
+        self,
+        left: np.ndarray,
+        right: np.ndarray,
+        closed: Literal["left", "right", "both", "neither"] = ...,
+        leaf_size: int = ...,
+    ): ...
+    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: ...

partial/pandas/_libs/tslibs/timedeltas.pyi

@@ -1,84 +1,83 @@
-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: Sequence[int], box: bool = ...) -> np.ndarray: ... 
+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: ...
-    @property
-    def resolution_string(self) -> str: ...
-    @property
-    def seconds(self) -> int: ...
-    max: Timedelta = ...
-    min: Timedelta = ...
-    resolution: 'Timedelta' = ...
+    def asm8(self) -> np.timedelta64: ...
+    # TODO: round/floor/ceil could return NaT?
+    def round(self, freq) -> Timedelta: ...
     def ceil(self, freq, **kwargs) -> Timedelta: ...
     def floor(self, freq, **kwargs) -> Timedelta: ...
+    @property
+    def resolution_string(self) -> str: ...
+    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: ...
     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) -> float: ...
-
-
-def delta_to_nanoseconds(delta: _Timedelta) -> int: ...
-def ints_to_pytimedelta(arr: Sequence[int], box: bool = ...) -> np.ndarray: ...