ENH: Imprpve interval_range and IntervalIndex (#351)

* ENH: IMrpve interval_range and IntervalIndex

* TYP: Improve typing accuracy

* TST: Add tests for IntervalIndex and interval_range

Co-authored-by: Kevin Sheppard <kevin.sheppard@gmail.com>

Author: bashtage

pandas-stubs/_typing.pyi

@@ -44,9 +44,7 @@ PandasScalar: TypeAlias = Union[
 ]
 # Scalar: TypeAlias = Union[PythonScalar, PandasScalar]
 
-DatetimeLike: TypeAlias = Union[
-    datetime.date, datetime.datetime, np.datetime64, Timestamp
-]
+DatetimeLike: TypeAlias = Union[datetime.datetime, np.datetime64, Timestamp]
 
 # dtypes
 NpDtype: TypeAlias = Union[

pandas-stubs/core/indexes/base.pyi

@@ -28,6 +28,7 @@ from pandas._typing import (
     Dtype,
     DtypeArg,
     DtypeObj,
+    FillnaOptions,
     HashableT,
     IndexT,
     Label,
@@ -155,7 +156,12 @@ class Index(IndexOpsMixin, PandasObject):
     def symmetric_difference(
         self, other: list[T1] | Index, result_name=..., sort=...
     ) -> Index: ...
-    def get_loc(self, key, tolerance=...): ...
+    def get_loc(
+        self,
+        key: Label,
+        method: FillnaOptions | Literal["nearest"] | None = ...,
+        tolerance=...,
+    ): ...
     def get_indexer(self, target, method=..., limit=..., tolerance=...): ...
     def reindex(self, target, method=..., level=..., limit=..., tolerance=...): ...
     def join(

pandas-stubs/core/indexes/interval.pyi

@@ -1,20 +1,50 @@
-from typing import Hashable
+import datetime as dt
+from typing import (
+    Any,
+    Hashable,
+    Literal,
+    Sequence,
+    Union,
+    overload,
+)
 
 import numpy as np
+import pandas as pd
+from pandas import Index
 from pandas.core.indexes.extension import ExtensionIndex
+from typing_extensions import TypeAlias
 
 from pandas._libs.interval import (
     Interval as Interval,
     IntervalMixin as IntervalMixin,
 )
+from pandas._libs.tslibs.offsets import DateOffset
 from pandas._typing import (
+    DatetimeLike,
     DtypeArg,
+    FillnaOptions,
     IntervalClosedType,
+    Label,
+    npt,
 )
 
 from pandas.core.dtypes.dtypes import IntervalDtype as IntervalDtype
 from pandas.core.dtypes.generic import ABCSeries
 
+_Edges: TypeAlias = Union[
+    Sequence[int],
+    Sequence[float],
+    Sequence[DatetimeLike],
+    npt.NDArray[np.int_],
+    npt.NDArray[np.float_],
+    npt.NDArray[np.datetime64],
+    pd.Series[int],
+    pd.Series[float],
+    pd.Series[pd.Timestamp],
+    pd.Int64Index,
+    pd.DatetimeIndex,
+]
+
 class IntervalIndex(IntervalMixin, ExtensionIndex):
     def __new__(
         cls,
@@ -28,7 +58,7 @@ class IntervalIndex(IntervalMixin, ExtensionIndex):
     @classmethod
     def from_breaks(
         cls,
-        breaks,
+        breaks: _Edges,
         closed: IntervalClosedType = ...,
         name: Hashable = ...,
         copy: bool = ...,
@@ -37,8 +67,8 @@ class IntervalIndex(IntervalMixin, ExtensionIndex):
     @classmethod
     def from_arrays(
         cls,
-        left,
-        right,
+        left: _Edges,
+        right: _Edges,
         closed: IntervalClosedType = ...,
         name: Hashable = ...,
         copy: bool = ...,
@@ -47,37 +77,67 @@ class IntervalIndex(IntervalMixin, ExtensionIndex):
     @classmethod
     def from_tuples(
         cls,
-        data,
+        data: Sequence[tuple[int, int]]
+        | Sequence[tuple[float, float]]
+        | Sequence[tuple[DatetimeLike, DatetimeLike]]
+        | npt.NDArray,
         closed: IntervalClosedType = ...,
         name: Hashable = ...,
         copy: bool = ...,
         dtype: IntervalDtype | None = ...,
     ) -> IntervalIndex: ...
+    def __contains__(self, key: Any) -> bool: ...
     def astype(self, dtype: DtypeArg, copy: bool = ...) -> IntervalIndex: ...
     @property
     def inferred_type(self) -> str: ...
     def memory_usage(self, deep: bool = ...) -> int: ...
     @property
     def is_overlapping(self) -> bool: ...
-    def get_loc(self, key, tolerance=...) -> int | slice | np.ndarray: ...
+    # Note: tolerance no effect. It is included in all get_loc so
+    # that signatures are consistent with base even though it is usually not used
+    def get_loc(
+        self,
+        key: Label,
+        method: FillnaOptions | Literal["nearest"] | None = ...,
+        tolerance=...,
+    ) -> int | slice | npt.NDArray[np.bool_]: ...
     def get_indexer(
         self,
-        targetArrayLike,
-        method: str | None = ...,
+        target: Index,
+        method: FillnaOptions | Literal["nearest"] | None = ...,
         limit: int | None = ...,
         tolerance=...,
-    ) -> np.ndarray: ...
+    ) -> npt.NDArray[np.intp]: ...
     def get_indexer_non_unique(
-        self, targetArrayLike
-    ) -> tuple[np.ndarray, np.ndarray]: ...
+        self, target: Index
+    ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ...
+    @property
+    def left(self) -> Index: ...
+    @property
+    def right(self) -> Index: ...
+    @property
+    def mid(self) -> Index: ...
+    @property
+    def length(self) -> Index: ...
     def get_value(self, series: ABCSeries, key): ...
     @property
     def is_all_dates(self) -> bool: ...
-    def __lt__(self, other): ...
-    def __le__(self, other): ...
-    def __gt__(self, other): ...
-    def __ge__(self, other): ...
 
+@overload
+def interval_range(
+    start: int | float | None = ...,
+    end: int | float | None = ...,
+    periods: int | None = ...,
+    freq: int | None = ...,
+    name: Hashable = ...,
+    closed: IntervalClosedType = ...,
+) -> IntervalIndex: ...
+@overload
 def interval_range(
-    start=..., end=..., periods=..., freq=..., name=..., closed: str = ...
-): ...
+    start: DatetimeLike | None = ...,
+    end: DatetimeLike | None = ...,
+    periods: int | None = ...,
+    freq: str | DateOffset | None = ...,
+    name: Hashable = ...,
+    closed: IntervalClosedType = ...,
+) -> IntervalIndex: ...