diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index 5913da7da..b749329ef 100644 --- a/pandas-stubs/_libs/interval.pyi +++ b/pandas-stubs/_libs/interval.pyi @@ -7,6 +7,7 @@ from typing import ( type_check_only, ) +import numpy as np from pandas import ( IntervalIndex, Series, @@ -65,9 +66,9 @@ class IntervalMixin: class Interval(IntervalMixin, Generic[_OrderableT]): @property - def left(self: Interval[_OrderableT]) -> _OrderableT: ... + def left(self) -> _OrderableT: ... @property - def right(self: Interval[_OrderableT]) -> _OrderableT: ... + def right(self) -> _OrderableT: ... @property def closed(self) -> IntervalClosedType: ... mid = _MidDescriptor() @@ -79,16 +80,10 @@ class Interval(IntervalMixin, Generic[_OrderableT]): closed: IntervalClosedType = ..., ) -> None: ... def __hash__(self) -> int: ... - # for __contains__, it seems that we have to separate out the 4 cases to make - # mypy happy @overload - def __contains__(self: Interval[Timestamp], key: Timestamp) -> bool: ... + def __contains__(self: Interval[int], key: float | np.floating) -> bool: ... @overload - def __contains__(self: Interval[Timedelta], key: Timedelta) -> bool: ... - @overload - def __contains__(self: Interval[int], key: float) -> bool: ... - @overload - def __contains__(self: Interval[float], key: float) -> bool: ... + def __contains__(self, key: _OrderableT) -> bool: ... @overload def __add__(self: Interval[Timestamp], y: Timedelta) -> Interval[Timestamp]: ... @overload diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index 36b992d70..c968845c8 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -7,6 +7,7 @@ from typing import ( Literal, TypeAlias, overload, + type_check_only, ) import numpy as np @@ -17,7 +18,11 @@ from pandas.core.indexes.extension import ExtensionIndex from pandas._libs.interval import ( Interval as Interval, IntervalMixin, + _OrderableScalarT, + _OrderableT, + _OrderableTimesT, ) +from pandas._libs.tslibs.timedeltas import Timedelta from pandas._typing import ( DatetimeLike, DtypeArg, @@ -58,6 +63,36 @@ _EdgesTimedelta: TypeAlias = ( _TimestampLike: TypeAlias = pd.Timestamp | np.datetime64 | dt.datetime _TimedeltaLike: TypeAlias = pd.Timedelta | np.timedelta64 | dt.timedelta +@type_check_only +class _LengthDescriptor: + @overload + def __get__( + self, + instance: IntervalIndex[Interval[_OrderableScalarT]], + owner: type[IntervalIndex], + ) -> Index[_OrderableScalarT]: ... + @overload + def __get__( + self, + instance: IntervalIndex[Interval[_OrderableTimesT]], + owner: type[IntervalIndex], + ) -> Index[Timedelta]: ... + +@type_check_only +class _MidDescriptor: + @overload + def __get__( + self, + instance: IntervalIndex[Interval[int]], + owner: type[IntervalIndex], + ) -> Index[float]: ... + @overload + def __get__( + self, + instance: IntervalIndex[Interval[_OrderableT]], + owner: type[IntervalIndex], + ) -> Index[_OrderableT]: ... + class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin): closed: IntervalClosedType @@ -216,16 +251,13 @@ class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin): def is_overlapping(self) -> bool: ... def get_loc(self, key: Label) -> int | slice | np_1darray_bool: ... @property - def left(self) -> Index: ... - @property - def right(self) -> Index: ... - @property - def mid(self) -> Index: ... + def left(self: IntervalIndex[Interval[_OrderableT]]) -> Index[_OrderableT]: ... @property - def length(self) -> Index: ... + def right(self: IntervalIndex[Interval[_OrderableT]]) -> Index[_OrderableT]: ... + mid = _MidDescriptor() + length = _LengthDescriptor() @overload # type: ignore[override] - # pyrefly: ignore # bad-override - def __getitem__( + def __getitem__( # pyrefly: ignore[bad-override] self, idx: ( slice diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 2d764df46..2db0c37c7 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -4,6 +4,7 @@ import datetime as dt import sys from typing import ( + TYPE_CHECKING, Any, cast, ) @@ -18,6 +19,7 @@ from pandas.core.indexes.base import Index from pandas.core.indexes.category import CategoricalIndex from pandas.core.indexes.datetimes import DatetimeIndex +import pytest from typing_extensions import ( Never, assert_type, @@ -847,6 +849,49 @@ def test_interval_index_tuples() -> None: ) +dt_l, dt_r = dt.datetime(2025, 12, 14), dt.datetime(2025, 12, 15) +td_l, td_r = dt.timedelta(seconds=1), dt.timedelta(seconds=2) + + +@pytest.mark.parametrize( + ("itv_idx", "typ_left", "typ_mid", "typ_length"), + [ + (pd.interval_range(0, 10), np.integer, np.floating, np.integer), + (pd.interval_range(0.0, 10), np.floating, np.floating, np.floating), + (pd.interval_range(dt_l, dt_r), pd.Timestamp, pd.Timestamp, pd.Timedelta), + (pd.interval_range(td_l, td_r, 2), pd.Timedelta, pd.Timedelta, pd.Timedelta), + ], +) +def test_interval_properties( + itv_idx: pd.IntervalIndex[Any], typ_left: type, typ_mid: type, typ_length: type +) -> None: + check(itv_idx.left, pd.Index, typ_left) + check(itv_idx.right, pd.Index, typ_left) + check(itv_idx.mid, pd.Index, typ_mid) + check(itv_idx.length, pd.Index, typ_length) + + if TYPE_CHECKING: + assert_type(pd.interval_range(0, 10).left, "pd.Index[int]") + assert_type(pd.interval_range(0, 10).right, "pd.Index[int]") + assert_type(pd.interval_range(0, 10).mid, "pd.Index[float]") + assert_type(pd.interval_range(0, 10).length, "pd.Index[int]") + + assert_type(pd.interval_range(0.0, 10).left, "pd.Index[float]") + assert_type(pd.interval_range(0.0, 10).right, "pd.Index[float]") + assert_type(pd.interval_range(0.0, 10).mid, "pd.Index[float]") + assert_type(pd.interval_range(0.0, 10).length, "pd.Index[float]") + + assert_type(pd.interval_range(dt_l, dt_r).left, "pd.Index[pd.Timestamp]") + assert_type(pd.interval_range(dt_l, dt_r).right, "pd.Index[pd.Timestamp]") + assert_type(pd.interval_range(dt_l, dt_r).mid, "pd.Index[pd.Timestamp]") + assert_type(pd.interval_range(dt_l, dt_r).length, "pd.Index[pd.Timedelta]") + + assert_type(pd.interval_range(td_l, td_r).left, "pd.Index[pd.Timedelta]") + assert_type(pd.interval_range(td_l, td_r).right, "pd.Index[pd.Timedelta]") + assert_type(pd.interval_range(td_l, td_r).mid, "pd.Index[pd.Timedelta]") + assert_type(pd.interval_range(td_l, td_r, 2).length, "pd.Index[pd.Timedelta]") + + def test_sorted_and_list() -> None: # GH 497 i1 = pd.Index([3, 2, 1])