ENH: Improve typing for Interval

.flake8

@@ -0,0 +1,3 @@
+[flake8]
+per-file-ignores =
+    tests/test_scalars.py: W503

pandas-stubs/_libs/interval.pyi

@@ -123,28 +123,35 @@ 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: ...
 
 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,7 @@ S1 = TypeVar(
     Timedelta,
     np.datetime64,
     Period,
+    Interval,
 )
 T1 = TypeVar(
     "T1", str, int, np.int64, np.uint64, np.float64, float, np.dtype[np.generic]

pandas-stubs/core/algorithms.pyi

@@ -1,5 +1,6 @@
 from typing import (
     Sequence,
+    TypeVar,
     overload,
 )
 
@@ -9,6 +10,7 @@ from pandas import (
     Categorical,
     CategoricalIndex,
     Index,
+    Interval,
     IntervalIndex,
     PeriodIndex,
     Series,
@@ -19,12 +21,21 @@ from pandas._typing import AnyArrayLike
 
 # 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
+
+_IntervalT = TypeVar(
+    "_IntervalT",
+    Interval[int],
+    Interval[float],
+    Interval[pd.Timestamp],
+    Interval[pd.Timedelta],
+)
+
 @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

pandas-stubs/core/indexes/interval.pyi

@@ -1,9 +1,11 @@
 import datetime as dt
 from typing import (
     Any,
+    Generic,
     Hashable,
     Literal,
     Sequence,
+    TypeVar,
     Union,
     overload,
 )
@@ -12,11 +14,15 @@ import numpy as np
 import pandas as pd
 from pandas import Index
 from pandas.core.indexes.extension import ExtensionIndex
+from pandas.core.series import (
+    TimedeltaSeries,
+    TimestampSeries,
+)
 from typing_extensions import TypeAlias
 
 from pandas._libs.interval import (
     Interval as Interval,
-    IntervalMixin as IntervalMixin,
+    IntervalMixin,
 )
 from pandas._libs.tslibs.offsets import DateOffset
 from pandas._typing import (
@@ -25,55 +31,144 @@ from pandas._typing import (
     FillnaOptions,
     IntervalClosedType,
     Label,
+    TimedeltaConvertibleTypes,
+    np_ndarray_bool,
     npt,
 )
 
 from pandas.core.dtypes.dtypes import IntervalDtype as IntervalDtype
 from pandas.core.dtypes.generic import ABCSeries
 
-_Edges: TypeAlias = Union[
+_EdgesInt: TypeAlias = Union[
     Sequence[int],
-    Sequence[float],
-    Sequence[DatetimeLike],
-    npt.NDArray[np.int_],
-    npt.NDArray[np.float_],
-    npt.NDArray[np.datetime64],
+    npt.NDArray[np.int64],
+    npt.NDArray[np.int32],
+    npt.NDArray[np.intp],
     pd.Series[int],
-    pd.Series[float],
-    pd.Series[pd.Timestamp],
     pd.Int64Index,
-    pd.DatetimeIndex,
 ]
+_EdgesFloat: TypeAlias = Union[
+    Sequence[float] | npt.NDArray[np.float64] | pd.Series[float] | pd.Float64Index,
+]
+_EdgesTimestamp: TypeAlias = Union[
+    Sequence[DatetimeLike]
+    | npt.NDArray[np.datetime64]
+    | pd.Series[pd.Timestamp]
+    | TimestampSeries
+    | pd.DatetimeIndex
+]
+_EdgesTimedelta: TypeAlias = Union[
+    Sequence[pd.Timedelta]
+    | npt.NDArray[np.timedelta64]
+    | pd.Series[pd.Timedelta]
+    | TimedeltaSeries
+    | pd.TimedeltaIndex
+]
+
+_IntervalT = TypeVar(
+    "_IntervalT",
+    Interval[int],
+    Interval[float],
+    Interval[pd.Timestamp],
+    Interval[pd.Timedelta],
+)
 
-class IntervalIndex(IntervalMixin, ExtensionIndex):
+class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[_IntervalT]):
     def __new__(
         cls,
-        data,
+        data: Sequence[_IntervalT],
         closed: IntervalClosedType = ...,
         dtype: IntervalDtype | None = ...,
         copy: bool = ...,
         name: Hashable = ...,
         verify_integrity: bool = ...,
-    ): ...
+    ) -> IntervalIndex[_IntervalT]: ...
+    # ignore[misc] here due to overlap, e.g., Sequence[int] and Sequence[float]
+    @overload
+    @classmethod
+    def from_breaks(  # type:ignore[misc]
+        cls,
+        breaks: _EdgesInt,
+        closed: IntervalClosedType = ...,
+        name: Hashable = ...,
+        copy: bool = ...,
+        dtype: IntervalDtype | None = ...,
+    ) -> IntervalIndex[Interval[int]]: ...
+    @overload
     @classmethod
     def from_breaks(
         cls,
-        breaks: _Edges,
+        breaks: _EdgesFloat,
         closed: IntervalClosedType = ...,
         name: Hashable = ...,
         copy: bool = ...,
         dtype: IntervalDtype | None = ...,
-    ) -> IntervalIndex: ...
+    ) -> IntervalIndex[Interval[float]]: ...
+    @overload
+    @classmethod
+    def from_breaks(
+        cls,
+        breaks: _EdgesTimestamp,
+        closed: IntervalClosedType = ...,
+        name: Hashable = ...,
+        copy: bool = ...,
+        dtype: IntervalDtype | None = ...,
+    ) -> IntervalIndex[Interval[pd.Timestamp]]: ...
+    @overload
+    @classmethod
+    def from_breaks(
+        cls,
+        breaks: _EdgesTimedelta,
+        closed: IntervalClosedType = ...,
+        name: Hashable = ...,
+        copy: bool = ...,
+        dtype: IntervalDtype | None = ...,
+    ) -> IntervalIndex[Interval[pd.Timedelta]]: ...
+    # ignore[misc] here due to overlap, e.g., Sequence[int] and Sequence[float]
+    @overload
+    @classmethod
+    def from_arrays(  # type:ignore[misc]
+        cls,
+        left: _EdgesInt,
+        right: _EdgesInt,
+        closed: IntervalClosedType = ...,
+        name: Hashable = ...,
+        copy: bool = ...,
+        dtype: IntervalDtype | None = ...,
+    ) -> IntervalIndex[Interval[int]]: ...
+    @overload
     @classmethod
     def from_arrays(
         cls,
-        left: _Edges,
-        right: _Edges,
+        left: _EdgesFloat,
+        right: _EdgesFloat,
         closed: IntervalClosedType = ...,
         name: Hashable = ...,
         copy: bool = ...,
         dtype: IntervalDtype | None = ...,
-    ) -> IntervalIndex: ...
+    ) -> IntervalIndex[Interval[float]]: ...
+    @overload
+    @classmethod
+    def from_arrays(
+        cls,
+        left: _EdgesTimestamp,
+        right: _EdgesTimestamp,
+        closed: IntervalClosedType = ...,
+        name: Hashable = ...,
+        copy: bool = ...,
+        dtype: IntervalDtype | None = ...,
+    ) -> IntervalIndex[Interval[pd.Timestamp]]: ...
+    @overload
+    @classmethod
+    def from_arrays(
+        cls,
+        left: _EdgesTimedelta,
+        right: _EdgesTimedelta,
+        closed: IntervalClosedType = ...,
+        name: Hashable = ...,
+        copy: bool = ...,
+        dtype: IntervalDtype | None = ...,
+    ) -> IntervalIndex[Interval[pd.Timedelta]]: ...
     @classmethod
     def from_tuples(
         cls,
@@ -122,6 +217,30 @@ class IntervalIndex(IntervalMixin, ExtensionIndex):
     def get_value(self, series: ABCSeries, key): ...
     @property
     def is_all_dates(self) -> bool: ...
+    @overload  # type: ignore[override]
+    def __gt__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ...  # type: ignore[misc]
+    @overload
+    def __gt__(self, other: object) -> bool: ...
+    @overload  # type: ignore[override]
+    def __ge__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ...  # type: ignore[misc]
+    @overload
+    def __ge__(self, other: object) -> bool: ...
+    @overload  # type: ignore[override]
+    def __le__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ...  # type: ignore[misc]
+    @overload
+    def __le__(self, other: object) -> bool: ...
+    @overload  # type: ignore[override]
+    def __lt__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ...  # type: ignore[misc]
+    @overload
+    def __lt__(self, other: object) -> bool: ...
+    @overload  # type: ignore[override]
+    def __eq__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ...  # type: ignore[misc]
+    @overload
+    def __eq__(self, other: object) -> bool: ...
+    @overload  # type: ignore[override]
+    def __ne__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ...  # type: ignore[misc]
+    @overload
+    def __ne__(self, other: object) -> bool: ...
 
 @overload
 def interval_range(

pandas-stubs/core/series.pyi

@@ -24,6 +24,7 @@ from matplotlib.axes import (
 )
 import numpy as np
 from pandas import (
+    Interval,
     Period,
     Timedelta,
     Timestamp,
@@ -43,6 +44,7 @@ from pandas.core.indexes.accessors import (
 )
 from pandas.core.indexes.base import Index
 from pandas.core.indexes.datetimes import DatetimeIndex
+from pandas.core.indexes.interval import IntervalIndex
 from pandas.core.indexes.period import PeriodIndex
 from pandas.core.indexes.timedeltas import TimedeltaIndex
 from pandas.core.indexing import (
@@ -196,6 +198,16 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
         fastpath: bool = ...,
     ) -> Series[Period]: ...
     @overload
+    def __new__(
+        cls,
+        data: IntervalIndex,
+        index: Axes | None = ...,
+        dtype=...,
+        name: Hashable | None = ...,
+        copy: bool = ...,
+        fastpath: bool = ...,
+    ) -> Series[Interval]: ...
+    @overload
     def __new__(
         cls,
         data: object | _ListLike | Series[S1] | dict[int, S1] | dict[_str, S1] | None,