Improve interval

Author: bashtage

pandas-stubs/_libs/interval.pyi

@@ -193,6 +193,8 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     def __ne__(self, other: Series[_OrderableT]) -> Series[bool]: ...  # type: ignore[misc]
     @overload
     def __ne__(self, other: object) -> Literal[True]: ...
+    @property
+    def is_empty(self) -> bool: ...
 
 class IntervalTree(IntervalMixin):
     def __init__(

pandas-stubs/core/arrays/interval.pyi

@@ -1,3 +1,11 @@
+import datetime as dt
+from typing import (
+    Generic,
+    Literal,
+    Sequence,
+    overload,
+)
+
 import numpy as np
 import pandas as pd
 from pandas import Index
@@ -9,21 +17,51 @@ from pandas._libs.interval import (
 )
 from pandas._typing import (
     Axis,
+    IntervalT,
     npt,
 )
 
 from pandas.core.dtypes.generic import ABCExtensionArray
 
-class IntervalArray(IntervalMixin, ExtensionArray):
+class IntervalArray(IntervalMixin, ExtensionArray, Generic[IntervalT]):
     ndim: int = ...
     can_hold_na: bool = ...
     def __new__(
-        cls, data, closed=..., dtype=..., copy: bool = ..., verify_integrity: bool = ...
-    ): ...
+        cls,
+        data: Sequence[IntervalT] | IntervalArray[IntervalT],
+        closed: Literal["left", "right", "both", "neither"] = ...,
+        dtype: pd.IntervalDtype = ...,
+        copy: bool = ...,
+        verify_integrity: bool = ...,
+    ) -> IntervalArray[IntervalT]: ...
+    @overload
     @classmethod
     def from_breaks(
-        cls, breaks, closed: str = ..., copy: bool = ..., dtype=...
-    ) -> IntervalArray: ...
+        cls, breaks: Sequence[int], closed: str = ..., copy: bool = ..., dtype=...
+    ) -> IntervalArray[Interval[int]]: ...
+    @overload
+    @classmethod
+    def from_breaks(
+        cls, breaks: Sequence[float], closed: str = ..., copy: bool = ..., dtype=...
+    ) -> IntervalArray[Interval[float]]: ...
+    @overload
+    @classmethod
+    def from_breaks(
+        cls,
+        breaks: Sequence[pd.Timestamp | np.datetime64 | dt.datetime],
+        closed: str = ...,
+        copy: bool = ...,
+        dtype=...,
+    ) -> IntervalArray[Interval[pd.Timestamp]]: ...
+    @overload
+    @classmethod
+    def from_breaks(
+        cls,
+        breaks: Sequence[pd.Timedelta | np.timedelta64 | dt.timedelta],
+        closed: str = ...,
+        copy: bool = ...,
+        dtype=...,
+    ) -> IntervalArray[Interval[pd.Timedelta]]: ...
     @classmethod
     def from_arrays(
         cls, left, right, closed: str = ..., copy: bool = ..., dtype=...
@@ -32,10 +70,10 @@ class IntervalArray(IntervalMixin, ExtensionArray):
     def from_tuples(
         cls, data, closed: str = ..., copy: bool = ..., dtype=...
     ) -> IntervalArray: ...
-    def __iter__(self): ...
+    def __iter__(self) -> IntervalT: ...
     def __len__(self) -> int: ...
-    def __getitem__(self, value): ...
-    def __setitem__(self, key, value) -> None: ...
+    def __getitem__(self, value: IntervalT): ...
+    def __setitem__(self, key: int, value: IntervalT) -> None: ...
     def __eq__(self, other): ...
     def __ne__(self, other): ...
     def fillna(self, value=..., method=..., limit=...): ...
@@ -61,7 +99,7 @@ class IntervalArray(IntervalMixin, ExtensionArray):
     def right(self) -> Index: ...
     @property
     def closed(self) -> str: ...
-    def set_closed(self, closed): ...
+    def set_closed(self, closed) -> IntervalArray[IntervalT]: ...
     @property
     def length(self) -> Index: ...
     @property
@@ -70,11 +108,11 @@ class IntervalArray(IntervalMixin, ExtensionArray):
     def is_non_overlapping_monotonic(self) -> bool: ...
     def __array__(self, dtype=...) -> np.ndarray: ...
     def __arrow_array__(self, type=...): ...
-    def to_tuples(self, na_tuple: bool = ...): ...
+    def to_tuples(self, na_tuple: bool = ...) -> npt.NDArray[np.object_]: ...
     def repeat(self, repeats, axis: Axis | None = ...): ...
     def contains(
         self, other: float | pd.Timestamp | pd.Timedelta
     ) -> npt.NDArray[np.bool_]: ...
-    def overlaps(self, other: Interval) -> bool: ...
+    def overlaps(self, other: Interval) -> npt.NDArray[np.bool_]: ...
     @property
     def is_empty(self) -> npt.NDArray[np.bool_]: ...