ENH: Improve typing for Interval (#391)

* ENH: Improve typing for Interval

* ENH: Further refinements of interval

* Attempt to make IntervalIndex generic

* ENH: Add more generic support

* ENH: Improve interval

* ENH: Improve from tuples

* CLN: Final clean ups

* CLN: Final clean ups

* CLN: Final fixes

* CLN: Final fixes

Author: bashtage

pandas-stubs/_libs/interval.pyi

@@ -1,18 +1,23 @@
 from typing import (
     Any,
     Generic,
+    Literal,
     TypeVar,
     overload,
 )
 
 import numpy as np
 from pandas import (
+    IntervalIndex,
+    Series,
     Timedelta,
     Timestamp,
 )
 
 from pandas._typing import (
     IntervalClosedType,
+    IntervalT,
+    np_ndarray_bool,
     npt,
 )
 
@@ -123,28 +128,73 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     @overload
     def __mul__(self: Interval[float], y: float) -> Interval[float]: ...
     @overload
+    def __mul__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
+    @overload
     def __rmul__(
         self: Interval[int], y: _OrderableScalarT
     ) -> Interval[_OrderableScalarT]: ...
     @overload
     def __rmul__(self: Interval[float], y: float) -> Interval[float]: ...
     @overload
-    def __truediv__(
-        self: Interval[int], y: _OrderableScalarT
-    ) -> Interval[_OrderableScalarT]: ...
+    def __rmul__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
+    @overload
+    def __truediv__(self: Interval[int], y: _OrderableScalarT) -> Interval[float]: ...
     @overload
     def __truediv__(self: Interval[float], y: float) -> Interval[float]: ...
     @overload
+    def __truediv__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
+    @overload
     def __floordiv__(
         self: Interval[int], y: _OrderableScalarT
     ) -> Interval[_OrderableScalarT]: ...
     @overload
     def __floordiv__(self: Interval[float], y: float) -> Interval[float]: ...
+    @overload
+    def __floordiv__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
+    @overload
     def overlaps(self: Interval[_OrderableT], other: Interval[_OrderableT]) -> bool: ...
-
-def intervals_to_interval_bounds(
-    intervals: np.ndarray, validate_closed: bool = ...
-) -> tuple[np.ndarray, np.ndarray, str]: ...
+    @overload
+    def overlaps(self: Interval[int], other: Interval[float]) -> bool: ...
+    @overload
+    def overlaps(self: Interval[float], other: Interval[int]) -> bool: ...
+    @overload
+    def __gt__(self, other: Interval[_OrderableT]) -> bool: ...
+    @overload
+    def __gt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    @overload
+    def __gt__(self, other: Series[_OrderableT]) -> Series[bool]: ...
+    @overload
+    def __lt__(self, other: Interval[_OrderableT]) -> bool: ...
+    @overload
+    def __lt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    @overload
+    def __lt__(self, other: Series[_OrderableT]) -> Series[bool]: ...
+    @overload
+    def __ge__(self, other: Interval[_OrderableT]) -> bool: ...
+    @overload
+    def __ge__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    @overload
+    def __ge__(self, other: Series[_OrderableT]) -> Series[bool]: ...
+    @overload
+    def __le__(self, other: Interval[_OrderableT]) -> bool: ...
+    @overload
+    def __le__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    @overload
+    def __eq__(self, other: Interval[_OrderableT]) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __eq__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...  # type: ignore[misc]
+    @overload
+    def __eq__(self, other: Series[_OrderableT]) -> Series[bool]: ...  # type: ignore[misc]
+    @overload
+    def __eq__(self, other: object) -> Literal[False]: ...
+    @overload
+    def __ne__(self, other: Interval[_OrderableT]) -> bool: ...  # type: ignore[misc]
+    @overload
+    def __ne__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...  # type: ignore[misc]
+    @overload
+    def __ne__(self, other: Series[_OrderableT]) -> Series[bool]: ...  # type: ignore[misc]
+    @overload
+    def __ne__(self, other: object) -> Literal[True]: ...
 
 class IntervalTree(IntervalMixin):
     def __init__(

pandas-stubs/_typing.pyi

@@ -26,6 +26,7 @@ from pandas.core.indexes.base import Index
 from pandas.core.series import Series
 from typing_extensions import TypeAlias
 
+from pandas._libs.interval import Interval
 from pandas._libs.tslibs import (
     Period,
     Timedelta,
@@ -196,6 +197,10 @@ S1 = TypeVar(
     Timedelta,
     np.datetime64,
     Period,
+    Interval[int],
+    Interval[float],
+    Interval[Timestamp],
+    Interval[Timedelta],
 )
 T1 = TypeVar(
     "T1", str, int, np.int64, np.uint64, np.float64, float, np.dtype[np.generic]
@@ -220,7 +225,13 @@ NDFrameT = TypeVar("NDFrameT", bound=NDFrame)
 IndexT = TypeVar("IndexT", bound=Index)
 
 # Interval closed type
-
+IntervalT = TypeVar(
+    "IntervalT",
+    Interval[int],
+    Interval[float],
+    Interval[Timestamp],
+    Interval[Timedelta],
+)
 IntervalClosedType: TypeAlias = Literal["left", "right", "both", "neither"]
 
 IgnoreRaiseCoerce: TypeAlias = Literal["ignore", "raise", "coerce"]

pandas-stubs/core/algorithms.pyi

@@ -14,16 +14,20 @@ from pandas import (
 )
 from pandas.api.extensions import ExtensionArray
 
-from pandas._typing import AnyArrayLike
+from pandas._typing import (
+    AnyArrayLike,
+    IntervalT,
+)
 
 # These are type: ignored because the Index types overlap due to inheritance but indices
 # with extension types return the same type while standard type return ndarray
+
 @overload
 def unique(values: PeriodIndex) -> PeriodIndex: ...  # type: ignore[misc]
 @overload
 def unique(values: CategoricalIndex) -> CategoricalIndex: ...  # type: ignore[misc]
 @overload
-def unique(values: IntervalIndex) -> IntervalIndex: ...  # type: ignore[misc]
+def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ...  # type: ignore[misc]
 @overload
 def unique(values: Index) -> np.ndarray: ...
 @overload