diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index df4e56304..69e79b7c8 100644 --- a/pandas-stubs/_libs/interval.pyi +++ b/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__( diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 42b4b5d40..22122ceaf 100644 --- a/pandas-stubs/_typing.pyi +++ b/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"] diff --git a/pandas-stubs/core/algorithms.pyi b/pandas-stubs/core/algorithms.pyi index 4582c3184..af56c0338 100644 --- a/pandas-stubs/core/algorithms.pyi +++ b/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 diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index dd5ef8510..5842ace9e 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -1,5 +1,6 @@ +import datetime as dt from typing import ( - Any, + Generic, Hashable, Literal, Sequence, @@ -11,81 +12,214 @@ 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.core.series import ( + TimedeltaSeries, + TimestampSeries, +) +from typing_extensions import ( + Never, + TypeAlias, +) from pandas._libs.interval import ( Interval as Interval, - IntervalMixin as IntervalMixin, + IntervalMixin, ) -from pandas._libs.tslibs.offsets import DateOffset +from pandas._libs.tslibs.offsets import BaseOffset from pandas._typing import ( DatetimeLike, DtypeArg, FillnaOptions, IntervalClosedType, + IntervalT, Label, + 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], + npt.NDArray[np.int64], + npt.NDArray[np.int32], + npt.NDArray[np.intp], + pd.Series[int], + pd.Int64Index, +] +_EdgesFloat: TypeAlias = Union[ Sequence[float], + npt.NDArray[np.float64], + pd.Series[float], + pd.Float64Index, +] +_EdgesTimestamp: TypeAlias = Union[ 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, + TimestampSeries, pd.DatetimeIndex, ] +_EdgesTimedelta: TypeAlias = Union[ + Sequence[pd.Timedelta], + npt.NDArray[np.timedelta64], + pd.Series[pd.Timedelta], + TimedeltaSeries, + pd.TimedeltaIndex, +] +_TimestampLike: TypeAlias = Union[pd.Timestamp, np.datetime64, dt.datetime] +_TimedeltaLike: TypeAlias = Union[pd.Timedelta, np.timedelta64, dt.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: _EdgesFloat, + closed: IntervalClosedType = ..., + name: Hashable = ..., + copy: bool = ..., + dtype: IntervalDtype | None = ..., + ) -> IntervalIndex[Interval[float]]: ... + @overload @classmethod def from_breaks( cls, - breaks: _Edges, + breaks: _EdgesTimestamp, closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., dtype: IntervalDtype | None = ..., - ) -> IntervalIndex: ... + ) -> 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: _EdgesFloat, + right: _EdgesFloat, + closed: IntervalClosedType = ..., + name: Hashable = ..., + copy: bool = ..., + dtype: IntervalDtype | None = ..., + ) -> IntervalIndex[Interval[float]]: ... + @overload @classmethod def from_arrays( cls, - left: _Edges, - right: _Edges, + left: _EdgesTimestamp, + right: _EdgesTimestamp, closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., dtype: IntervalDtype | None = ..., - ) -> IntervalIndex: ... + ) -> 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]]: ... + @overload @classmethod def from_tuples( cls, - data: Sequence[tuple[int, int]] - | Sequence[tuple[float, float]] - | Sequence[tuple[DatetimeLike, DatetimeLike]] - | npt.NDArray, + data: Sequence[tuple[int, int]], closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., dtype: IntervalDtype | None = ..., - ) -> IntervalIndex: ... - def __contains__(self, key: Any) -> bool: ... + ) -> IntervalIndex[pd.Interval[int]]: ... + # Ignore misc here due to intentional overlap between int and float + @overload + @classmethod + def from_tuples( + cls, + data: Sequence[tuple[float, float]], + closed: IntervalClosedType = ..., + name: Hashable = ..., + copy: bool = ..., + dtype: IntervalDtype | None = ..., + ) -> IntervalIndex[pd.Interval[float]]: ... + @overload + @classmethod + def from_tuples( + cls, + data: Sequence[ + tuple[pd.Timestamp, pd.Timestamp] + | tuple[dt.datetime, dt.datetime] + | tuple[np.datetime64, np.datetime64] + ], + closed: IntervalClosedType = ..., + name: Hashable = ..., + copy: bool = ..., + dtype: IntervalDtype | None = ..., + ) -> IntervalIndex[pd.Interval[pd.Timestamp]]: ... + @overload + @classmethod + def from_tuples( + cls, + data: Sequence[ + tuple[pd.Timedelta, pd.Timedelta] + | tuple[dt.timedelta, dt.timedelta] + | tuple[np.timedelta64, np.timedelta64] + ], + closed: IntervalClosedType = ..., + name: Hashable = ..., + copy: bool = ..., + dtype: IntervalDtype | None = ..., + ) -> IntervalIndex[pd.Interval[pd.Timedelta]]: ... + @overload + def __contains__(self, key: IntervalT) -> bool: ... # type: ignore[misc] + @overload + def __contains__(self, key: object) -> Literal[False]: ... def astype(self, dtype: DtypeArg, copy: bool = ...) -> IntervalIndex: ... @property def inferred_type(self) -> str: ... @@ -121,22 +255,164 @@ class IntervalIndex(IntervalMixin, ExtensionIndex): def get_value(self, series: ABCSeries, key): ... @property def is_all_dates(self) -> bool: ... + # override is due to additional types for comparison + # misc is due to overlap with object below + @overload # type: ignore[override] + def __gt__( # type: ignore[misc] + self, other: IntervalT | IntervalIndex[IntervalT] + ) -> np_ndarray_bool: ... + @overload + def __gt__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[misc] + @overload + def __gt__(self, other: object) -> Never: ... + @overload # type: ignore[override] + def __ge__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] + @overload + def __ge__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[misc] + @overload + def __ge__(self, other: object) -> Never: ... + @overload # type: ignore[override] + def __le__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] + @overload + def __le__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[misc] + @overload + def __le__(self, other: object) -> Never: ... + @overload # type: ignore[override] + def __lt__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] + @overload + def __lt__(self, other: pd.Series[IntervalT]) -> bool: ... # type: ignore[misc] + @overload + def __lt__(self, other: object) -> Never: ... + @overload # type: ignore[override] + def __eq__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] + @overload + def __eq__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[misc] + @overload + def __eq__(self, other: object) -> Literal[False]: ... + @overload # type: ignore[override] + def __ne__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] + @overload + def __ne__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[misc] + @overload + def __ne__(self, other: object) -> Literal[True]: ... + +# misc here because int and float overlap but interval has distinct types +# int gets hit first and so the correct type is returned +@overload +def interval_range( # type: ignore[misc] + start: int = ..., + end: int = ..., + periods: int | None = ..., + freq: int | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[int]]: ... +# Overlaps since int is a subclass of float +@overload +def interval_range( # pyright: reportOverlappingOverload=false + start: int, + *, + end: None = ..., + periods: int | None = ..., + freq: int | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[int]]: ... +@overload +def interval_range( # pyright: reportOverlappingOverload=false + *, + start: None = ..., + end: int, + periods: int | None = ..., + freq: int | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[int]]: ... +@overload +def interval_range( + start: float = ..., + end: float = ..., + periods: int | None = ..., + freq: int | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[float]]: ... @overload def interval_range( - start: int | float | None = ..., - end: int | float | None = ..., + start: float, + *, + end: None = ..., periods: int | None = ..., freq: int | None = ..., name: Hashable = ..., closed: IntervalClosedType = ..., -) -> IntervalIndex: ... +) -> IntervalIndex[Interval[float]]: ... +@overload +def interval_range( + *, + start: None = ..., + end: float, + periods: int | None = ..., + freq: int | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[float]]: ... +@overload +def interval_range( + start: _TimestampLike, + end: _TimestampLike = ..., + periods: int | None = ..., + freq: str | BaseOffset | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[pd.Timestamp]]: ... +@overload +def interval_range( + *, + start: None = ..., + end: _TimestampLike, + periods: int | None = ..., + freq: str | BaseOffset | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[pd.Timestamp]]: ... +@overload +def interval_range( + start: _TimestampLike, + *, + end: None = ..., + periods: int | None = ..., + freq: str | BaseOffset | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[pd.Timestamp]]: ... +@overload +def interval_range( + start: _TimedeltaLike, + end: _TimedeltaLike = ..., + periods: int | None = ..., + freq: str | BaseOffset | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[pd.Timedelta]]: ... +@overload +def interval_range( + *, + start: None = ..., + end: _TimedeltaLike, + periods: int | None = ..., + freq: str | BaseOffset | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[pd.Timedelta]]: ... @overload def interval_range( - start: DatetimeLike | None = ..., - end: DatetimeLike | None = ..., + start: _TimedeltaLike, + *, + end: None = ..., periods: int | None = ..., - freq: str | DateOffset | None = ..., + freq: str | BaseOffset | None = ..., name: Hashable = ..., closed: IntervalClosedType = ..., -) -> IntervalIndex: ... +) -> IntervalIndex[Interval[pd.Timedelta]]: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 21ea1d1b9..76f5fc239 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -43,6 +43,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 ( @@ -91,6 +92,7 @@ from pandas._typing import ( IgnoreRaise, IndexingInt, IntervalClosedType, + IntervalT, JoinHow, JsonSeriesOrient, Level, @@ -210,6 +212,16 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): fastpath: bool = ..., ) -> TimedeltaSeries: ... @overload + def __new__( + cls, + data: IntervalIndex[IntervalT], + index: Axes | None = ..., + dtype=..., + name: Hashable | None = ..., + copy: bool = ..., + fastpath: bool = ..., + ) -> Series[IntervalT]: ... + @overload def __new__( cls, data: object | _ListLike | Series[S1] | dict[int, S1] | dict[_str, S1] | None, diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 18e8e0203..cd4d5c9ca 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -191,45 +191,59 @@ def test_range_index_union(): def test_interval_range(): - check(assert_type(pd.interval_range(0, 10), pd.IntervalIndex), pd.IntervalIndex) + check( + assert_type(pd.interval_range(0, 10), "pd.IntervalIndex[pd.Interval[int]]"), + pd.IntervalIndex, + pd.Interval, + ) check( assert_type( - pd.interval_range(0, 10, name="something", closed="both"), pd.IntervalIndex + pd.interval_range(0, 10, name="something", closed="both"), + "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, + ) + check( + assert_type(pd.interval_range(0.0, 10), "pd.IntervalIndex[pd.Interval[float]]"), + pd.IntervalIndex, + pd.Interval, ) - check(assert_type(pd.interval_range(0.0, 10), pd.IntervalIndex), pd.IntervalIndex) check( assert_type( pd.interval_range(dt.datetime(2000, 1, 1), dt.datetime(2010, 1, 1), 5), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( pd.interval_range( np.datetime64("2000-01-01"), np.datetime64("2020-01-01"), 5 ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( pd.interval_range(pd.Timestamp(2000, 1, 1), pd.Timestamp(2010, 1, 1), 5), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( pd.interval_range( pd.Timestamp(2000, 1, 1), pd.Timestamp(2010, 1, 1), freq="1M" ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -238,87 +252,158 @@ def test_interval_range(): pd.Timestamp(2010, 1, 1), freq=pd.DateOffset(months=2), ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( pd.interval_range(pd.Timestamp(2000, 1, 1), dt.datetime(2010, 1, 1), 5), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) - -def test_interval_index_breaks(): check( - assert_type(pd.IntervalIndex.from_breaks([1, 2, 3, 4]), pd.IntervalIndex), + assert_type( + pd.interval_range(pd.Timedelta("1D"), pd.Timedelta("10D")), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + check( + assert_type( + pd.interval_range(end=pd.Timedelta("10D"), periods=10, freq="D"), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", + ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( - pd.IntervalIndex.from_breaks([1.0, 2.0, 3.0, 4.0]), pd.IntervalIndex + pd.interval_range(start=pd.Timedelta("1D"), periods=10, freq="D"), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", ), pd.IntervalIndex, + pd.Interval, ) + + +def test_interval_index_breaks(): check( assert_type( - pd.IntervalIndex.from_breaks(np.array([1, 2, 3, 4])), pd.IntervalIndex + pd.IntervalIndex.from_breaks([1, 2, 3, 4]), + "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( - pd.IntervalIndex.from_breaks(np.array([1.0, 2.0, 3.0, 4.0])), - pd.IntervalIndex, + pd.IntervalIndex.from_breaks([1.0, 2.0, 3.0, 4.0]), + "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( pd.IntervalIndex.from_breaks( - np.array( - [ - np.datetime64("2000-01-01"), - np.datetime64("2001-01-01"), - np.datetime64("2002-01-01"), - np.datetime64("2003-01-01"), - ] - ) + [pd.Timestamp(2000, 1, 1), pd.Timestamp(2000, 1, 2)] ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( - pd.IntervalIndex.from_breaks(pd.Series([1, 2, 3, 4])), pd.IntervalIndex + pd.IntervalIndex.from_breaks([pd.Timedelta(1, "D"), pd.Timedelta(2, "D")]), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", ), pd.IntervalIndex, + pd.Interval, ) + check( assert_type( - pd.IntervalIndex.from_breaks(pd.Series([1.0, 2.0, 3.0, 4.0])), - pd.IntervalIndex, + pd.IntervalIndex.from_breaks(np.array([1, 2, 3, 4], dtype=np.int64)), + "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( pd.IntervalIndex.from_breaks( - pd.Series( - [ - pd.Timestamp(2000, 1, 1), - pd.Timestamp(2001, 1, 1), - pd.Timestamp(2002, 1, 1), - pd.Timestamp(2003, 1, 1), - ] - ) + np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float64) ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, + ) + np_ndarray_dt64 = np.array( + [ + np.datetime64("2000-01-01"), + np.datetime64("2001-01-01"), + np.datetime64("2002-01-01"), + np.datetime64("2003-01-01"), + ], + dtype=np.datetime64, + ) + check( + assert_type( + pd.IntervalIndex.from_breaks(np_ndarray_dt64), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + np_ndarray_td64 = np.array( + [ + np.timedelta64(1, "D"), + np.timedelta64(2, "D"), + np.timedelta64(3, "D"), + np.timedelta64(4, "D"), + ], + dtype=np.timedelta64, + ) + check( + assert_type( + pd.IntervalIndex.from_breaks(np_ndarray_td64), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + check( + assert_type( + pd.IntervalIndex.from_breaks(pd.Series([1, 2, 3, 4], dtype=int)), + "pd.IntervalIndex[pd.Interval[int]]", + ), + pd.IntervalIndex, + ) + pd_series_float = pd.Series([1.0, 2.0, 3.0, 4.0], dtype=float) + check( + assert_type( + pd.IntervalIndex.from_breaks(pd_series_float), + "pd.IntervalIndex[pd.Interval[float]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + timestamp_series = pd.Series(pd.date_range("2000-01-01", "2003-01-01", freq="D")) + check( + assert_type( + pd.IntervalIndex.from_breaks(timestamp_series), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", + ), + pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -330,110 +415,108 @@ def test_interval_index_breaks(): dt.datetime(2003, 1, 1), ] ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) def test_interval_index_arrays(): check( assert_type( - pd.IntervalIndex.from_arrays([1, 2, 3, 4], [2, 3, 4, 5]), pd.IntervalIndex + pd.IntervalIndex.from_arrays([1, 2, 3, 4], [2, 3, 4, 5]), + "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( pd.IntervalIndex.from_arrays([1.0, 2.0, 3.0, 4.0], [2.0, 3.0, 4.0, 5.0]), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( pd.IntervalIndex.from_arrays( - np.array([1, 2, 3, 4]), np.array([2, 3, 4, 5]) + np.array([1, 2, 3, 4], dtype=np.int64), + np.array([2, 3, 4, 5], dtype=np.int64), ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( pd.IntervalIndex.from_arrays( - np.array([1.0, 2.0, 3.0, 4.0]), np.array([2.0, 3.0, 4.0, 5.0]) + np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float64), + np.array([2.0, 3.0, 4.0, 5.0], dtype=np.float64), ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, + ) + left_dt64_arr: npt.NDArray[np.datetime64] = np.array( + [ + np.datetime64("2000-01-01"), + np.datetime64("2001-01-01"), + np.datetime64("2002-01-01"), + np.datetime64("2003-01-01"), + ], + dtype="datetime64[ns]", + ) + right_dt_arr: npt.NDArray[np.datetime64] = np.array( + [ + np.datetime64("2001-01-01"), + np.datetime64("2002-01-01"), + np.datetime64("2003-01-01"), + np.datetime64("2004-01-01"), + ], + dtype="datetime64[ns]", ) check( assert_type( - pd.IntervalIndex.from_arrays( - np.array( - [ - np.datetime64("2000-01-01"), - np.datetime64("2001-01-01"), - np.datetime64("2002-01-01"), - np.datetime64("2003-01-01"), - ] - ), - np.array( - [ - np.datetime64("2001-01-01"), - np.datetime64("2002-01-01"), - np.datetime64("2003-01-01"), - np.datetime64("2004-01-01"), - ] - ), - ), - pd.IntervalIndex, + pd.IntervalIndex.from_arrays(left_dt64_arr, right_dt_arr), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( pd.IntervalIndex.from_arrays( - pd.Series([1, 2, 3, 4]), pd.Series([2, 3, 4, 5]) + pd.Series([1, 2, 3, 4], dtype=int), pd.Series([2, 3, 4, 5], dtype=int) ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, ) + series_float_left = pd.Series([1.0, 2.0, 3.0, 4.0], dtype=float) + series_float_right = pd.Series([2.0, 3.0, 4.0, 5.0], dtype=float) check( assert_type( - pd.IntervalIndex.from_arrays( - pd.Series([1.0, 2.0, 3.0, 4.0]), pd.Series([2.0, 3.0, 4.0, 5.0]) - ), - pd.IntervalIndex, + pd.IntervalIndex.from_arrays(series_float_left, series_float_right), + "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, ) + left_s_ts = pd.Series(pd.date_range("2000-01-01", "2003-01-01", freq="Y")) + right_s_ts = pd.Series(pd.date_range("2001-01-01", "2004-01-01", freq="Y")) check( assert_type( - pd.IntervalIndex.from_arrays( - pd.Series( - [ - pd.Timestamp(2000, 1, 1), - pd.Timestamp(2001, 1, 1), - pd.Timestamp(2002, 1, 1), - pd.Timestamp(2003, 1, 1), - ] - ), - pd.Series( - [ - pd.Timestamp(2001, 1, 1), - pd.Timestamp(2002, 1, 1), - pd.Timestamp(2003, 1, 1), - pd.Timestamp(2004, 1, 1), - ] - ), - ), - pd.IntervalIndex, + pd.IntervalIndex.from_arrays(left_s_ts, right_s_ts), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -451,22 +534,29 @@ def test_interval_index_arrays(): dt.datetime(2004, 1, 1), ], ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) def test_interval_index_tuples(): check( - assert_type(pd.IntervalIndex.from_tuples([(1, 2), (2, 3)]), pd.IntervalIndex), + assert_type( + pd.IntervalIndex.from_tuples([(1, 2), (2, 3)]), + "pd.IntervalIndex[pd.Interval[int]]", + ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( - pd.IntervalIndex.from_tuples([(1.0, 2.0), (2.0, 3.0)]), pd.IntervalIndex + pd.IntervalIndex.from_tuples([(1.0, 2.0), (2.0, 3.0)]), + "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -476,7 +566,73 @@ def test_interval_index_tuples(): (pd.Timestamp(2001, 1, 1), pd.Timestamp(2002, 1, 1)), ] ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + check( + assert_type( + pd.IntervalIndex.from_tuples( + [ + (dt.datetime(2000, 1, 1), dt.datetime(2001, 1, 1)), + (dt.datetime(2001, 1, 1), dt.datetime(2002, 1, 1)), + ] + ), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + check( + assert_type( + pd.IntervalIndex.from_tuples( + [ + (np.datetime64("2000-01-01"), np.datetime64("2001-01-01")), + (np.datetime64("2001-01-01"), np.datetime64("2002-01-01")), + ] + ), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + check( + assert_type( + pd.IntervalIndex.from_tuples( + [ + (pd.Timedelta(1, "D"), pd.Timedelta(2, "D")), + (pd.Timedelta(2, "D"), pd.Timedelta(3, "D")), + ] + ), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + check( + assert_type( + pd.IntervalIndex.from_tuples( + [ + (dt.timedelta(days=1), dt.timedelta(days=2)), + (dt.timedelta(days=2), dt.timedelta(days=3)), + ] + ), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + check( + assert_type( + pd.IntervalIndex.from_tuples( + [ + (np.timedelta64(1, "D"), np.timedelta64(2, "D")), + (np.timedelta64(2, "D"), np.timedelta64(3, "D")), + ] + ), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", ), pd.IntervalIndex, + pd.Interval, ) diff --git a/tests/test_pandas.py b/tests/test_pandas.py index 456f94947..e295ea9dc 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -775,7 +775,10 @@ def test_index_unqiue() -> None: check(assert_type(pd.unique(ui), np.ndarray), np.ndarray) check(assert_type(pd.unique(tdi), np.ndarray), np.ndarray) check(assert_type(pd.unique(mi), np.ndarray), np.ndarray) - check(assert_type(pd.unique(interval_i), pd.IntervalIndex), pd.IntervalIndex) + check( + assert_type(pd.unique(interval_i), "pd.IntervalIndex[pd.Interval[int]]"), + pd.IntervalIndex, + ) def test_cut() -> None: diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 5c788909d..53dbd28ba 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -51,6 +51,339 @@ OffsetSeries: TypeAlias = pd.Series +def test_interval() -> None: + interval_i = pd.Interval(0, 1, closed="left") + interval_f = pd.Interval(0.0, 1.0, closed="right") + interval_ts = pd.Interval( + pd.Timestamp("2017-01-01"), pd.Timestamp("2017-01-02"), closed="both" + ) + interval_td = pd.Interval( + pd.Timedelta("1 days"), pd.Timedelta("2 days"), closed="neither" + ) + + check(assert_type(interval_i, "pd.Interval[int]"), pd.Interval, int) + check(assert_type(interval_f, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(interval_ts, "pd.Interval[pd.Timestamp]"), pd.Interval, pd.Timestamp + ) + check( + assert_type(interval_td, "pd.Interval[pd.Timedelta]"), pd.Interval, pd.Timedelta + ) + + check( + assert_type(interval_i.closed, Literal["left", "right", "both", "neither"]), str + ) + check(assert_type(interval_i.closed_left, bool), bool) + check(assert_type(interval_i.closed_right, bool), bool) + check(assert_type(interval_i.is_empty, bool), bool) + check(assert_type(interval_i.left, int), int) + check(assert_type(interval_i.length, int), int) + check(assert_type(interval_i.mid, float), float) + check(assert_type(interval_i.open_left, bool), bool) + check(assert_type(interval_i.open_right, bool), bool) + check(assert_type(interval_i.right, int), int) + + check( + assert_type(interval_f.closed, Literal["left", "right", "both", "neither"]), str + ) + check(assert_type(interval_f.closed_left, bool), bool) + check(assert_type(interval_f.closed_right, bool), bool) + check(assert_type(interval_f.is_empty, bool), bool) + check(assert_type(interval_f.left, float), float) + check(assert_type(interval_f.length, float), float) + check(assert_type(interval_f.mid, float), float) + check(assert_type(interval_f.open_left, bool), bool) + check(assert_type(interval_f.open_right, bool), bool) + check(assert_type(interval_f.right, float), float) + + check( + assert_type(interval_ts.closed, Literal["left", "right", "both", "neither"]), + str, + ) + check(assert_type(interval_ts.closed_left, bool), bool) + check(assert_type(interval_ts.closed_right, bool), bool) + check(assert_type(interval_ts.is_empty, bool), bool) + check(assert_type(interval_ts.left, pd.Timestamp), pd.Timestamp) + check(assert_type(interval_ts.length, pd.Timedelta), pd.Timedelta) + check(assert_type(interval_ts.mid, pd.Timestamp), pd.Timestamp) + check(assert_type(interval_ts.open_left, bool), bool) + check(assert_type(interval_ts.open_right, bool), bool) + check(assert_type(interval_ts.right, pd.Timestamp), pd.Timestamp) + + check( + assert_type(interval_td.closed, Literal["left", "right", "both", "neither"]), + str, + ) + check(assert_type(interval_td.closed_left, bool), bool) + check(assert_type(interval_td.closed_right, bool), bool) + check(assert_type(interval_td.is_empty, bool), bool) + check(assert_type(interval_td.left, pd.Timedelta), pd.Timedelta) + check(assert_type(interval_td.length, pd.Timedelta), pd.Timedelta) + check(assert_type(interval_td.mid, pd.Timedelta), pd.Timedelta) + check(assert_type(interval_td.open_left, bool), bool) + check(assert_type(interval_td.open_right, bool), bool) + check(assert_type(interval_td.right, pd.Timedelta), pd.Timedelta) + + check( + assert_type(interval_i.overlaps(pd.Interval(0.5, 1.5, closed="left")), bool), + bool, + ) + check( + assert_type(interval_i.overlaps(pd.Interval(2, 3, closed="left")), bool), bool + ) + + check( + assert_type(interval_f.overlaps(pd.Interval(0.5, 1.5, closed="left")), bool), + bool, + ) + check( + assert_type(interval_f.overlaps(pd.Interval(2, 3, closed="left")), bool), bool + ) + ts1 = pd.Timestamp(year=2017, month=1, day=1) + ts2 = pd.Timestamp(year=2017, month=1, day=2) + check( + assert_type(interval_ts.overlaps(pd.Interval(ts1, ts2, closed="left")), bool), + bool, + ) + td1 = pd.Timedelta(days=1) + td2 = pd.Timedelta(days=3) + check( + assert_type(interval_td.overlaps(pd.Interval(td1, td2, closed="left")), bool), + bool, + ) + + +def test_interval_math() -> None: + interval_i = pd.Interval(0, 1, closed="left") + interval_f = pd.Interval(0.0, 1.0, closed="right") + interval_ts = pd.Interval( + pd.Timestamp("2017-01-01"), pd.Timestamp("2017-01-02"), closed="both" + ) + interval_td = pd.Interval( + pd.Timedelta("1 days"), pd.Timedelta("2 days"), closed="neither" + ) + + check(assert_type(interval_i * 3, "pd.Interval[int]"), pd.Interval, int) + check(assert_type(interval_f * 3, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(interval_td * 3, "pd.Interval[pd.Timedelta]"), + pd.Interval, + pd.Timedelta, + ) + + check(assert_type(interval_i * 3.5, "pd.Interval[float]"), pd.Interval, float) + check(assert_type(interval_f * 3.5, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(interval_td * 3.5, "pd.Interval[pd.Timedelta]"), + pd.Interval, + pd.Timedelta, + ) + + check(assert_type(3 * interval_i, "pd.Interval[int]"), pd.Interval, int) + check(assert_type(3 * interval_f, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(3 * interval_td, "pd.Interval[pd.Timedelta]"), + pd.Interval, + pd.Timedelta, + ) + + check(assert_type(3.5 * interval_i, "pd.Interval[float]"), pd.Interval, float) + check(assert_type(3.5 * interval_f, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(3.5 * interval_td, "pd.Interval[pd.Timedelta]"), + pd.Interval, + pd.Timedelta, + ) + + check(assert_type(interval_i / 3, "pd.Interval[float]"), pd.Interval, float) + check(assert_type(interval_f / 3, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(interval_td / 3, "pd.Interval[pd.Timedelta]"), + pd.Interval, + pd.Timedelta, + ) + + check(assert_type(interval_i / 3.5, "pd.Interval[float]"), pd.Interval, float) + check(assert_type(interval_f / 3.5, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(interval_td / 3.5, "pd.Interval[pd.Timedelta]"), + pd.Interval, + pd.Timedelta, + ) + + check(assert_type(interval_i // 3, "pd.Interval[int]"), pd.Interval, int) + check(assert_type(interval_f // 3, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(interval_td // 3, "pd.Interval[pd.Timedelta]"), + pd.Interval, + pd.Timedelta, + ) + + check(assert_type(interval_i // 3.5, "pd.Interval[float]"), pd.Interval, float) + check(assert_type(interval_f // 3.5, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(interval_td // 3.5, "pd.Interval[pd.Timedelta]"), + pd.Interval, + pd.Timedelta, + ) + + # Subtraction + check(assert_type(interval_i - 1, "pd.Interval[int]"), pd.Interval, int) + check(assert_type(interval_f - 1, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(interval_ts - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), + pd.Interval, + pd.Timestamp, + ) + check( + assert_type(interval_td - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), + pd.Interval, + pd.Timedelta, + ) + + check(assert_type(interval_i - 1.5, "pd.Interval[float]"), pd.Interval, float) + check(assert_type(interval_f - 1.5, "pd.Interval[float]"), pd.Interval, float) + check( + assert_type(interval_ts - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), + pd.Interval, + ) + check( + assert_type(interval_td - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), + pd.Interval, + ) + + # Addition + check(assert_type(interval_i + 1, "pd.Interval[int]"), pd.Interval) + check(assert_type(1 + interval_i, "pd.Interval[int]"), pd.Interval) + check(assert_type(interval_f + 1, "pd.Interval[float]"), pd.Interval) + check(assert_type(1 + interval_f, "pd.Interval[float]"), pd.Interval) + check( + assert_type(interval_ts + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), + pd.Interval, + ) + check( + assert_type(pd.Timedelta(days=1) + interval_ts, "pd.Interval[pd.Timestamp]"), + pd.Interval, + ) + check( + assert_type(interval_td + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), + pd.Interval, + ) + check( + assert_type(pd.Timedelta(days=1) + interval_td, "pd.Interval[pd.Timedelta]"), + pd.Interval, + ) + + check(assert_type(interval_i + 1.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(1.5 + interval_i, "pd.Interval[float]"), pd.Interval) + check(assert_type(interval_f + 1.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(1.5 + interval_f, "pd.Interval[float]"), pd.Interval) + check( + assert_type(interval_ts + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), + pd.Interval, + ) + check( + assert_type(pd.Timedelta(days=1) + interval_ts, "pd.Interval[pd.Timestamp]"), + pd.Interval, + ) + check( + assert_type(interval_td + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), + pd.Interval, + ) + check( + assert_type(pd.Timedelta(days=1) + interval_td, "pd.Interval[pd.Timedelta]"), + pd.Interval, + ) + + +def test_interval_cmp(): + interval_i = pd.Interval(0, 1, closed="left") + interval_f = pd.Interval(0.0, 1.0, closed="right") + interval_ts = pd.Interval( + pd.Timestamp("2017-01-01"), pd.Timestamp("2017-01-02"), closed="both" + ) + interval_td = pd.Interval( + pd.Timedelta("1 days"), pd.Timedelta("2 days"), closed="neither" + ) + + check(assert_type(0.5 in interval_i, bool), bool) + check(assert_type(1 in interval_i, bool), bool) + check(assert_type(1 in interval_f, bool), bool) + check(assert_type(pd.Timestamp("2000-1-1") in interval_ts, bool), bool) + check(assert_type(pd.Timedelta(days=1) in interval_td, bool), bool) + + interval_index_int = pd.IntervalIndex([interval_i]) + check( + assert_type(interval_index_int >= interval_i, np_ndarray_bool), + np.ndarray, + np.bool_, + ) + check( + assert_type(interval_index_int < interval_i, np_ndarray_bool), + np.ndarray, + np.bool_, + ) + check( + assert_type(interval_index_int <= interval_i, np_ndarray_bool), + np.ndarray, + np.bool_, + ) + check( + assert_type(interval_index_int > interval_i, np_ndarray_bool), + np.ndarray, + np.bool_, + ) + + check( + assert_type(interval_i >= interval_index_int, np_ndarray_bool), + np.ndarray, + np.bool_, + ) + check( + assert_type(interval_i < interval_index_int, np_ndarray_bool), + np.ndarray, + np.bool_, + ) + check( + assert_type(interval_i <= interval_index_int, np_ndarray_bool), + np.ndarray, + np.bool_, + ) + check( + assert_type(interval_i > interval_index_int, np_ndarray_bool), + np.ndarray, + np.bool_, + ) + + check( + assert_type(interval_index_int == interval_i, np_ndarray_bool), + np.ndarray, + np.bool_, + ) + check( + assert_type(interval_index_int != interval_i, np_ndarray_bool), + np.ndarray, + np.bool_, + ) + + check( + assert_type( + interval_i == interval_index_int, + np_ndarray_bool, + ), + np.ndarray, + np.bool_, + ) + check( + assert_type( + interval_i != interval_index_int, + np_ndarray_bool, + ), + np.ndarray, + np.bool_, + ) + + def test_timedelta_construction() -> None: check(assert_type(pd.Timedelta(1, "H"), pd.Timedelta), pd.Timedelta)