From 977f98f25b4ec0510791470c8b5c962c615496ff Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Sun, 16 Oct 2022 13:58:38 +0100 Subject: [PATCH 01/10] ENH: Improve typing for Interval --- pandas-stubs/_libs/interval.pyi | 17 +++- tests/test_scalars.py | 156 ++++++++++++++++++++++++++++++++ 2 files changed, 169 insertions(+), 4 deletions(-) create mode 100644 tests/test_scalars.py diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index df4e56304..c726cef62 100644 --- a/pandas-stubs/_libs/interval.pyi +++ b/pandas-stubs/_libs/interval.pyi @@ -123,28 +123,37 @@ 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 __rmul__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ... + @overload def __truediv__( self: Interval[int], y: _OrderableScalarT ) -> Interval[_OrderableScalarT]: ... @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__( diff --git a/tests/test_scalars.py b/tests/test_scalars.py new file mode 100644 index 000000000..b9c9f1986 --- /dev/null +++ b/tests/test_scalars.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from typing import Literal + +import pandas as pd +from typing_extensions import assert_type + +from tests import check + + +def test_interval() -> None: + i0 = pd.Interval(0, 1, closed="left") + i1 = pd.Interval(0.0, 1.0, closed="right") + i2 = pd.Interval( + pd.Timestamp("2017-01-01"), pd.Timestamp("2017-01-02"), closed="both" + ) + i3 = pd.Interval(pd.Timedelta("1 days"), pd.Timedelta("2 days"), closed="neither") + check(assert_type(i0, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1, "pd.Interval[float]"), pd.Interval) + check(assert_type(i2, "pd.Interval[pd.Timestamp]"), pd.Interval) + check(assert_type(i3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0.closed, Literal["left", "right", "both", "neither"]), str) + check(assert_type(i0.closed_left, bool), bool) + check(assert_type(i0.closed_right, bool), bool) + check(assert_type(i0.is_empty, bool), bool) + check(assert_type(i0.left, int), int) + check(assert_type(i0.length, int), int) + check(assert_type(i0.mid, float), float) + check(assert_type(i0.open_left, bool), bool) + check(assert_type(i0.open_right, bool), bool) + check(assert_type(i0.right, int), int) + + check(assert_type(i1.closed, Literal["left", "right", "both", "neither"]), str) + check(assert_type(i1.closed_left, bool), bool) + check(assert_type(i1.closed_right, bool), bool) + check(assert_type(i1.is_empty, bool), bool) + check(assert_type(i1.left, float), float) + check(assert_type(i1.length, float), float) + check(assert_type(i1.mid, float), float) + check(assert_type(i1.open_left, bool), bool) + check(assert_type(i1.open_right, bool), bool) + check(assert_type(i1.right, float), float) + + check(assert_type(i2.closed, Literal["left", "right", "both", "neither"]), str) + check(assert_type(i2.closed_left, bool), bool) + check(assert_type(i2.closed_right, bool), bool) + check(assert_type(i2.is_empty, bool), bool) + check(assert_type(i2.left, pd.Timestamp), pd.Timestamp) + check(assert_type(i2.length, pd.Timedelta), pd.Timedelta) + check(assert_type(i2.mid, pd.Timestamp), pd.Timestamp) + check(assert_type(i2.open_left, bool), bool) + check(assert_type(i2.open_right, bool), bool) + check(assert_type(i2.right, pd.Timestamp), pd.Timestamp) + + check(assert_type(i3.closed, Literal["left", "right", "both", "neither"]), str) + check(assert_type(i3.closed_left, bool), bool) + check(assert_type(i3.closed_right, bool), bool) + check(assert_type(i3.is_empty, bool), bool) + check(assert_type(i3.left, pd.Timedelta), pd.Timedelta) + check(assert_type(i3.length, pd.Timedelta), pd.Timedelta) + check(assert_type(i3.mid, pd.Timedelta), pd.Timedelta) + check(assert_type(i3.open_left, bool), bool) + check(assert_type(i3.open_right, bool), bool) + check(assert_type(i3.right, pd.Timedelta), pd.Timedelta) + + check(assert_type(i0.overlaps(pd.Interval(0.5, 1.5, closed="left")), bool), bool) + check(assert_type(i0.overlaps(pd.Interval(2, 3, closed="left")), bool), bool) + + check(assert_type(i1.overlaps(pd.Interval(0.5, 1.5, closed="left")), bool), bool) + check(assert_type(i1.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(i2.overlaps(pd.Interval(ts1, ts2, closed="left")), bool), bool) + td1 = pd.Timedelta(days=1) + td2 = pd.Timedelta(days=3) + check(assert_type(i3.overlaps(pd.Interval(td1, td2, closed="left")), bool), bool) + + check(assert_type(i0 * 3, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1 * 3, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 * 3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 * 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i1 * 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 * 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(3 * i0, "pd.Interval[int]"), pd.Interval) + check(assert_type(3 * i1, "pd.Interval[float]"), pd.Interval) + check(assert_type(3 * i3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(3.5 * i0, "pd.Interval[float]"), pd.Interval) + check(assert_type(3.5 * i1, "pd.Interval[float]"), pd.Interval) + check(assert_type(3.5 * i3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 / 3, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1 / 3, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 / 3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 / 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i1 / 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 / 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 // 3, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1 // 3, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 // 3, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 // 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i1 // 3.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i3 // 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval) + + check(assert_type(i0 - 1, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1 - 1, "pd.Interval[float]"), pd.Interval) + check( + assert_type(i2 - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + ) + check( + assert_type(i3 - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + ) + + check(assert_type(i0 - 1.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i1 - 1.5, "pd.Interval[float]"), pd.Interval) + check( + assert_type(i2 - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + ) + check( + assert_type(i3 - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + ) + + check(assert_type(i0 + 1, "pd.Interval[int]"), pd.Interval) + check(assert_type(i1 + 1, "pd.Interval[float]"), pd.Interval) + check( + assert_type(i2 + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + ) + check( + assert_type(i3 + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + ) + + check(assert_type(i0 + 1.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(i1 + 1.5, "pd.Interval[float]"), pd.Interval) + check( + assert_type(i2 + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + ) + check( + assert_type(i3 + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + ) + + check(assert_type(0.5 in i0, bool), bool) + check(assert_type(1 in i0, bool), bool) + check(assert_type(1 in i1, bool), bool) + check(assert_type(pd.Timestamp("2000-1-1") in i2, bool), bool) + check(assert_type(pd.Timedelta(days=1) in i3, bool), bool) + + check(assert_type(hash(i0), int), int) + check(assert_type(hash(i1), int), int) + check(assert_type(hash(i2), int), int) + check(assert_type(hash(i3), int), int) From dbb1735d6c6460dc8eb228c0cd1f90d8f064b4e0 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Fri, 21 Oct 2022 19:14:20 +0100 Subject: [PATCH 02/10] ENH: Further refinements of interval --- .flake8 | 3 + pandas-stubs/_libs/interval.pyi | 4 +- pandas-stubs/_typing.pyi | 2 + pandas-stubs/core/indexes/interval.pyi | 25 ++ pandas-stubs/core/series.pyi | 12 + tests/test_scalars.py | 349 +++++++++++++++++-------- 6 files changed, 283 insertions(+), 112 deletions(-) create mode 100644 .flake8 diff --git a/.flake8 b/.flake8 new file mode 100644 index 000000000..fe708d67b --- /dev/null +++ b/.flake8 @@ -0,0 +1,3 @@ +[flake8] +per-file-ignores = + tests/test_scalars.py: W503 \ No newline at end of file diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index c726cef62..7f105fb62 100644 --- a/pandas-stubs/_libs/interval.pyi +++ b/pandas-stubs/_libs/interval.pyi @@ -133,9 +133,7 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __rmul__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ... @overload - def __truediv__( - self: Interval[int], y: _OrderableScalarT - ) -> Interval[_OrderableScalarT]: ... + def __truediv__(self: Interval[int], y: _OrderableScalarT) -> Interval[float]: ... @overload def __truediv__(self: Interval[float], y: float) -> Interval[float]: ... @overload diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 42b4b5d40..ff972f781 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,7 @@ S1 = TypeVar( Timedelta, np.datetime64, Period, + Interval, ) T1 = TypeVar( "T1", str, int, np.int64, np.uint64, np.float64, float, np.dtype[np.generic] diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index 339f99053..f25f321e2 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -25,6 +25,7 @@ from pandas._typing import ( FillnaOptions, IntervalClosedType, Label, + np_ndarray_bool, npt, ) @@ -122,6 +123,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( diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 0a12a9a1b..db7bf22b3 100644 --- a/pandas-stubs/core/series.pyi +++ b/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, diff --git a/tests/test_scalars.py b/tests/test_scalars.py index b9c9f1986..d62d8eec2 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -2,6 +2,7 @@ from typing import Literal +import numpy as np import pandas as pd from typing_extensions import assert_type @@ -9,148 +10,278 @@ def test_interval() -> None: - i0 = pd.Interval(0, 1, closed="left") - i1 = pd.Interval(0.0, 1.0, closed="right") - i2 = pd.Interval( + 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" ) - i3 = pd.Interval(pd.Timedelta("1 days"), pd.Timedelta("2 days"), closed="neither") - check(assert_type(i0, "pd.Interval[int]"), pd.Interval) - check(assert_type(i1, "pd.Interval[float]"), pd.Interval) - check(assert_type(i2, "pd.Interval[pd.Timestamp]"), pd.Interval) - check(assert_type(i3, "pd.Interval[pd.Timedelta]"), pd.Interval) - - check(assert_type(i0.closed, Literal["left", "right", "both", "neither"]), str) - check(assert_type(i0.closed_left, bool), bool) - check(assert_type(i0.closed_right, bool), bool) - check(assert_type(i0.is_empty, bool), bool) - check(assert_type(i0.left, int), int) - check(assert_type(i0.length, int), int) - check(assert_type(i0.mid, float), float) - check(assert_type(i0.open_left, bool), bool) - check(assert_type(i0.open_right, bool), bool) - check(assert_type(i0.right, int), int) - - check(assert_type(i1.closed, Literal["left", "right", "both", "neither"]), str) - check(assert_type(i1.closed_left, bool), bool) - check(assert_type(i1.closed_right, bool), bool) - check(assert_type(i1.is_empty, bool), bool) - check(assert_type(i1.left, float), float) - check(assert_type(i1.length, float), float) - check(assert_type(i1.mid, float), float) - check(assert_type(i1.open_left, bool), bool) - check(assert_type(i1.open_right, bool), bool) - check(assert_type(i1.right, float), float) - - check(assert_type(i2.closed, Literal["left", "right", "both", "neither"]), str) - check(assert_type(i2.closed_left, bool), bool) - check(assert_type(i2.closed_right, bool), bool) - check(assert_type(i2.is_empty, bool), bool) - check(assert_type(i2.left, pd.Timestamp), pd.Timestamp) - check(assert_type(i2.length, pd.Timedelta), pd.Timedelta) - check(assert_type(i2.mid, pd.Timestamp), pd.Timestamp) - check(assert_type(i2.open_left, bool), bool) - check(assert_type(i2.open_right, bool), bool) - check(assert_type(i2.right, pd.Timestamp), pd.Timestamp) - - check(assert_type(i3.closed, Literal["left", "right", "both", "neither"]), str) - check(assert_type(i3.closed_left, bool), bool) - check(assert_type(i3.closed_right, bool), bool) - check(assert_type(i3.is_empty, bool), bool) - check(assert_type(i3.left, pd.Timedelta), pd.Timedelta) - check(assert_type(i3.length, pd.Timedelta), pd.Timedelta) - check(assert_type(i3.mid, pd.Timedelta), pd.Timedelta) - check(assert_type(i3.open_left, bool), bool) - check(assert_type(i3.open_right, bool), bool) - check(assert_type(i3.right, pd.Timedelta), pd.Timedelta) - - check(assert_type(i0.overlaps(pd.Interval(0.5, 1.5, closed="left")), bool), bool) - check(assert_type(i0.overlaps(pd.Interval(2, 3, closed="left")), bool), bool) - - check(assert_type(i1.overlaps(pd.Interval(0.5, 1.5, closed="left")), bool), bool) - check(assert_type(i1.overlaps(pd.Interval(2, 3, closed="left")), bool), bool) + 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(i2.overlaps(pd.Interval(ts1, ts2, closed="left")), bool), bool) + 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(i3.overlaps(pd.Interval(td1, td2, closed="left")), bool), bool) + check( + assert_type(interval_td.overlaps(pd.Interval(td1, td2, closed="left")), bool), + bool, + ) - check(assert_type(i0 * 3, "pd.Interval[int]"), pd.Interval) - check(assert_type(i1 * 3, "pd.Interval[float]"), pd.Interval) - check(assert_type(i3 * 3, "pd.Interval[pd.Timedelta]"), pd.Interval) + 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(i0 * 3.5, "pd.Interval[float]"), pd.Interval) - check(assert_type(i1 * 3.5, "pd.Interval[float]"), pd.Interval) - check(assert_type(i3 * 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval) + 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 * i0, "pd.Interval[int]"), pd.Interval) - check(assert_type(3 * i1, "pd.Interval[float]"), pd.Interval) - check(assert_type(3 * i3, "pd.Interval[pd.Timedelta]"), pd.Interval) + 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(3.5 * i0, "pd.Interval[float]"), pd.Interval) - check(assert_type(3.5 * i1, "pd.Interval[float]"), pd.Interval) - check(assert_type(3.5 * i3, "pd.Interval[pd.Timedelta]"), pd.Interval) + 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(i0 / 3, "pd.Interval[int]"), pd.Interval) - check(assert_type(i1 / 3, "pd.Interval[float]"), pd.Interval) - check(assert_type(i3 / 3, "pd.Interval[pd.Timedelta]"), pd.Interval) + 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(i0 / 3.5, "pd.Interval[float]"), pd.Interval) - check(assert_type(i1 / 3.5, "pd.Interval[float]"), pd.Interval) - check(assert_type(i3 / 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval) + 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(i0 // 3, "pd.Interval[int]"), pd.Interval) - check(assert_type(i1 // 3, "pd.Interval[float]"), pd.Interval) - check(assert_type(i3 // 3, "pd.Interval[pd.Timedelta]"), pd.Interval) + 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(i0 // 3.5, "pd.Interval[float]"), pd.Interval) - check(assert_type(i1 // 3.5, "pd.Interval[float]"), pd.Interval) - check(assert_type(i3 // 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval) + 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(i0 - 1, "pd.Interval[int]"), pd.Interval) - check(assert_type(i1 - 1, "pd.Interval[float]"), pd.Interval) + 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(i2 - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + assert_type(interval_ts - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), + pd.Interval, ) check( - assert_type(i3 - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + assert_type(interval_td - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), + pd.Interval, ) - check(assert_type(i0 - 1.5, "pd.Interval[float]"), pd.Interval) - check(assert_type(i1 - 1.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(interval_i + 1, "pd.Interval[int]"), pd.Interval) + check(assert_type(interval_f + 1, "pd.Interval[float]"), pd.Interval) check( - assert_type(i2 - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + assert_type(interval_ts + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), + pd.Interval, ) check( - assert_type(i3 - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + assert_type(interval_td + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), + pd.Interval, ) - check(assert_type(i0 + 1, "pd.Interval[int]"), pd.Interval) - check(assert_type(i1 + 1, "pd.Interval[float]"), pd.Interval) + check(assert_type(interval_i + 1.5, "pd.Interval[float]"), pd.Interval) + check(assert_type(interval_f + 1.5, "pd.Interval[float]"), pd.Interval) check( - assert_type(i2 + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + assert_type(interval_ts + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), + pd.Interval, ) check( - assert_type(i3 + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + assert_type(interval_td + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), + pd.Interval, ) - check(assert_type(i0 + 1.5, "pd.Interval[float]"), pd.Interval) - check(assert_type(i1 + 1.5, "pd.Interval[float]"), pd.Interval) + 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) + + check(assert_type(hash(interval_i), int), int) + check(assert_type(hash(interval_f), int), int) + check(assert_type(hash(interval_ts), int), int) + check(assert_type(hash(interval_td), int), int) + + interval_index_int = pd.IntervalIndex([interval_i]) + interval_series_int = pd.Series(interval_index_int) + + check(interval_series_int >= interval_i, pd.Series, bool) + check(interval_series_int < interval_i, pd.Series, bool) + check(interval_series_int <= interval_i, pd.Series, bool) + check(interval_series_int > interval_i, pd.Series, bool) + + check(interval_i >= interval_series_int, pd.Series, bool) + check(interval_i < interval_series_int, pd.Series, bool) + check(interval_i <= interval_series_int, pd.Series, bool) + check(interval_i > interval_series_int, pd.Series, bool) + + check(interval_series_int == interval_i, pd.Series, bool) + check(interval_series_int != interval_i, pd.Series, bool) + check( - assert_type(i2 + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval + interval_i + == interval_series_int, # pyright: ignore[reportUnnecessaryComparison] + pd.Series, + bool, ) check( - assert_type(i3 + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval + interval_i + != interval_series_int, # pyright: ignore[reportUnnecessaryComparison] + pd.Series, + bool, ) - check(assert_type(0.5 in i0, bool), bool) - check(assert_type(1 in i0, bool), bool) - check(assert_type(1 in i1, bool), bool) - check(assert_type(pd.Timestamp("2000-1-1") in i2, bool), bool) - check(assert_type(pd.Timedelta(days=1) in i3, bool), bool) + check(interval_index_int >= interval_i, np.ndarray, np.bool_) + check(interval_index_int < interval_i, np.ndarray, np.bool_) + check(interval_index_int <= interval_i, np.ndarray, np.bool_) + check(interval_index_int > interval_i, np.ndarray, np.bool_) + + check(interval_i >= interval_index_int, np.ndarray, np.bool_) + check(interval_i < interval_index_int, np.ndarray, np.bool_) + check(interval_i <= interval_index_int, np.ndarray, np.bool_) + check(interval_i > interval_index_int, np.ndarray, np.bool_) + + check(interval_index_int == interval_i, np.ndarray, np.bool_) + check(interval_index_int != interval_i, np.ndarray, np.bool_) - check(assert_type(hash(i0), int), int) - check(assert_type(hash(i1), int), int) - check(assert_type(hash(i2), int), int) - check(assert_type(hash(i3), int), int) + check( + interval_i + == interval_index_int, # pyright: ignore[reportUnnecessaryComparison] + np.ndarray, + np.bool_, + ) + check( + interval_i + != interval_index_int, # pyright: ignore[reportUnnecessaryComparison] + np.ndarray, + np.bool_, + ) From f1805f5f6f38a97e3c528e1c363ebfa5b2e06c73 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Sun, 23 Oct 2022 09:07:54 +0100 Subject: [PATCH 03/10] Attempt to make IntervalIndex generic --- pandas-stubs/core/algorithms.pyi | 13 ++- pandas-stubs/core/indexes/interval.pyi | 130 +++++++++++++++++++++---- tests/test_indexes.py | 106 +++++++++++++------- 3 files changed, 196 insertions(+), 53 deletions(-) diff --git a/pandas-stubs/core/algorithms.pyi b/pandas-stubs/core/algorithms.pyi index 4ac2621a1..14ca6d0fd 100644 --- a/pandas-stubs/core/algorithms.pyi +++ b/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 diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index f25f321e2..ca1423564 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/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,6 +31,7 @@ from pandas._typing import ( FillnaOptions, IntervalClosedType, Label, + TimedeltaConvertibleTypes, np_ndarray_bool, npt, ) @@ -32,49 +39,136 @@ from pandas._typing import ( 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 ] -class IntervalIndex(IntervalMixin, ExtensionIndex): +_IntervalT = TypeVar( + "_IntervalT", + Interval[int], + Interval[float], + Interval[pd.Timestamp], + Interval[pd.Timedelta], +) + +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, diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 18e8e0203..8de8c9010 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -253,70 +253,105 @@ def test_interval_range(): def test_interval_index_breaks(): check( - assert_type(pd.IntervalIndex.from_breaks([1, 2, 3, 4]), pd.IntervalIndex), + assert_type( + pd.IntervalIndex.from_breaks([1, 2, 3, 4]), + "pd.IntervalIndex[pd.Interval[int]]", + ), pd.IntervalIndex, ) check( assert_type( - pd.IntervalIndex.from_breaks([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, ) check( assert_type( - pd.IntervalIndex.from_breaks(np.array([1, 2, 3, 4])), pd.IntervalIndex + pd.IntervalIndex.from_breaks( + [pd.Timestamp(2000, 1, 1), pd.Timestamp(2000, 1, 2)] + ), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, ) check( assert_type( - pd.IntervalIndex.from_breaks(np.array([1.0, 2.0, 3.0, 4.0])), - pd.IntervalIndex, + pd.IntervalIndex.from_breaks([pd.Timedelta(1, "D"), pd.Timedelta(2, "D")]), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", + ), + pd.IntervalIndex, + ) + + check( + assert_type( + pd.IntervalIndex.from_breaks(np.array([1, 2, 3, 4], dtype=np.int64)), + "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, ) 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"), - ] - ) + np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float64) ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, ) + 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(pd.Series([1, 2, 3, 4])), pd.IntervalIndex + pd.IntervalIndex.from_breaks(np_ndarray_dt64), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, ) + np_ndarray_td64 = np.array( + [ + np.timedelta64(1), + np.timedelta64(2), + np.timedelta64(3), + np.timedelta64(4), + ], + dtype=np.timedelta64, + ) check( assert_type( - pd.IntervalIndex.from_breaks(pd.Series([1.0, 2.0, 3.0, 4.0])), - pd.IntervalIndex, + pd.IntervalIndex.from_breaks(np_ndarray_td64), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", + ), + pd.IntervalIndex, + ) + 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, ) 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), - ] - ) + pd.Series(pd.date_range("2000-01-01", "2003-01-01", freq="D")) ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, ) @@ -330,7 +365,7 @@ def test_interval_index_breaks(): dt.datetime(2003, 1, 1), ] ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, ) @@ -339,32 +374,35 @@ def test_interval_index_breaks(): 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, ) 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, ) 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, ) 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, ) @@ -451,7 +489,7 @@ def test_interval_index_arrays(): dt.datetime(2004, 1, 1), ], ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, ) From 66846d3223685fb7bb2cb521c30f4ac73bb8e496 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Mon, 24 Oct 2022 14:14:53 +0100 Subject: [PATCH 04/10] ENH: Add more generic support --- pandas-stubs/_libs/interval.pyi | 89 ++++++++++++ pandas-stubs/_typing.pyi | 13 +- pandas-stubs/core/algorithms.pyi | 17 +-- pandas-stubs/core/indexes/interval.pyi | 191 +++++++++++++++++++------ pandas-stubs/core/series.pyi | 6 +- tests/test_indexes.py | 58 ++++---- tests/test_pandas.py | 5 +- tests/test_scalars.py | 95 ++++++------ 8 files changed, 347 insertions(+), 127 deletions(-) diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index 7f105fb62..d044697f3 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, ) @@ -152,6 +157,90 @@ class Interval(IntervalMixin, Generic[_OrderableT]): 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: ... # type: ignore[misc] + @overload + def __gt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] + # @overload + # def __gt__(self:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... + # @overload + # def __gt__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... + # @overload + # def __gt__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... + @overload + def __gt__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[misc] + @overload + def __gt__(self, other: object) -> Literal[False]: ... + @overload + def __lt__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[misc] + @overload + def __lt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] + # @overload + # def __lt__(self:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... + # @overload + # def __lt__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... + # @overload + # def __lt__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... + @overload + def __lt__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[misc] + @overload + def __lt__(self, other: object) -> Literal[False]: ... + @overload + def __ge__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[misc] + @overload + def __ge__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] + # @overload + # def __ge__(self:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... + # @overload + # def __ge__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... + # @overload + # def __ge__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... + @overload + def __ge__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[misc] + @overload + def __ge__(self, other: object) -> Literal[False]: ... + @overload + def __le__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[misc] + @overload + def __le__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] + # @overload + # def __le__(self:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... + # @overload + # def __le__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... + # @overload + # def __le__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... + # @overload + # def __le__(self, other: Series[_OrderableT]) -> Series[bool]: ... + @overload + def __le__(self, other: object) -> Literal[False]: ... + @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:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... + # @overload + # def __eq__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... + # @overload + # def __eq__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... + @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:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... + # @overload + # def __ne__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... + # @overload + # def __ne__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... + @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 ff972f781..22122ceaf 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -197,7 +197,10 @@ S1 = TypeVar( Timedelta, np.datetime64, Period, - Interval, + Interval[int], + Interval[float], + Interval[Timestamp], + Interval[Timedelta], ) T1 = TypeVar( "T1", str, int, np.int64, np.uint64, np.float64, float, np.dtype[np.generic] @@ -222,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 14ca6d0fd..eb8c45608 100644 --- a/pandas-stubs/core/algorithms.pyi +++ b/pandas-stubs/core/algorithms.pyi @@ -1,6 +1,5 @@ from typing import ( Sequence, - TypeVar, overload, ) @@ -10,32 +9,26 @@ from pandas import ( Categorical, CategoricalIndex, Index, - Interval, IntervalIndex, PeriodIndex, Series, ) 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 -_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[_IntervalT]) -> IntervalIndex[_IntervalT]: ... # 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 ca1423564..1e71d5dc8 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -5,7 +5,6 @@ from typing import ( Hashable, Literal, Sequence, - TypeVar, Union, overload, ) @@ -24,14 +23,14 @@ from pandas._libs.interval import ( Interval as Interval, 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, - TimedeltaConvertibleTypes, np_ndarray_bool, npt, ) @@ -48,41 +47,38 @@ _EdgesInt: TypeAlias = Union[ pd.Int64Index, ] _EdgesFloat: TypeAlias = Union[ - Sequence[float] | npt.NDArray[np.float64] | pd.Series[float] | pd.Float64Index, + 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 + 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 + 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] -_IntervalT = TypeVar( - "_IntervalT", - Interval[int], - Interval[float], - Interval[pd.Timestamp], - Interval[pd.Timedelta], -) - -class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[_IntervalT]): +class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]): def __new__( cls, - data: Sequence[_IntervalT], + data: Sequence[IntervalT], closed: IntervalClosedType = ..., dtype: IntervalDtype | None = ..., copy: bool = ..., name: Hashable = ..., verify_integrity: bool = ..., - ) -> IntervalIndex[_IntervalT]: ... + ) -> IntervalIndex[IntervalT]: ... # ignore[misc] here due to overlap, e.g., Sequence[int] and Sequence[float] @overload @classmethod @@ -217,46 +213,153 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[_IntervalT]): 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__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ... # type: ignore[misc] + def __gt__( # type: ignore[misc] + self, other: IntervalT | IntervalIndex[IntervalT] + ) -> np_ndarray_bool: ... @overload - def __gt__(self, other: object) -> bool: ... + def __gt__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[misc] + @overload + def __gt__(self, other: object) -> Literal[False]: ... @overload # type: ignore[override] - def __ge__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ... # type: ignore[misc] + 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) -> bool: ... + def __ge__(self, other: object) -> Literal[False]: ... @overload # type: ignore[override] - def __le__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ... # type: ignore[misc] + def __le__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] @overload - def __le__(self, other: object) -> bool: ... + def __le__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[misc] + @overload + def __le__(self, other: object) -> Literal[False]: ... @overload # type: ignore[override] - def __lt__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ... # type: ignore[misc] + 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) -> bool: ... + def __lt__(self, other: object) -> Literal[False]: ... @overload # type: ignore[override] - def __eq__(self, other: Interval | IntervalIndex) -> np_ndarray_bool: ... # type: ignore[misc] + def __eq__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] @overload - def __eq__(self, other: object) -> bool: ... + 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: Interval | IntervalIndex) -> np_ndarray_bool: ... # type: ignore[misc] + 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) -> bool: ... + 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]]: ... @overload def interval_range( - start: int | float | None = ..., - end: int | float | None = ..., + start: int, + *, + end: None = ..., periods: int | None = ..., freq: int | None = ..., name: Hashable = ..., closed: IntervalClosedType = ..., -) -> IntervalIndex: ... +) -> IntervalIndex[Interval[int]]: ... +@overload +def interval_range( + *, + 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: float, + *, + end: None = ..., + periods: int | None = ..., + freq: int | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> 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: None = ..., + end: _TimedeltaLike, + periods: int | None = ..., + freq: str | BaseOffset | None = ..., + name: Hashable = ..., + closed: IntervalClosedType = ..., +) -> IntervalIndex[Interval[pd.Timestamp]]: ... @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 f9a5bbd96..af43cf244 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -24,7 +24,6 @@ from matplotlib.axes import ( ) import numpy as np from pandas import ( - Interval, Period, Timedelta, Timestamp, @@ -89,6 +88,7 @@ from pandas._typing import ( IgnoreRaise, IndexingInt, IntervalClosedType, + IntervalT, JoinHow, JsonSeriesOrient, Level, @@ -210,13 +210,13 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def __new__( cls, - data: IntervalIndex, + data: IntervalIndex[IntervalT], index: Axes | None = ..., dtype=..., name: Hashable | None = ..., copy: bool = ..., fastpath: bool = ..., - ) -> Series[Interval]: ... + ) -> Series[IntervalT]: ... @overload def __new__( cls, diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 8de8c9010..53508ae1a 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -191,18 +191,25 @@ 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, + ) 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, ) - check(assert_type(pd.interval_range(0.0, 10), pd.IntervalIndex), pd.IntervalIndex) + check( + assert_type(pd.interval_range(0.0, 10), "pd.IntervalIndex[pd.Interval[float]]"), + 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, ) @@ -211,14 +218,14 @@ def test_interval_range(): pd.interval_range( np.datetime64("2000-01-01"), np.datetime64("2020-01-01"), 5 ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, ) 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, ) @@ -227,7 +234,7 @@ def test_interval_range(): 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, ) @@ -238,14 +245,14 @@ def test_interval_range(): pd.Timestamp(2010, 1, 1), freq=pd.DateOffset(months=2), ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, ) 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, ) @@ -317,10 +324,10 @@ def test_interval_index_breaks(): ) np_ndarray_td64 = np.array( [ - np.timedelta64(1), - np.timedelta64(2), - np.timedelta64(3), - np.timedelta64(4), + np.timedelta64(1, "D"), + np.timedelta64(2, "D"), + np.timedelta64(3, "D"), + np.timedelta64(4, "D"), ], dtype=np.timedelta64, ) @@ -346,11 +353,10 @@ def test_interval_index_breaks(): ), pd.IntervalIndex, ) + timestamp_series = pd.Series(pd.date_range("2000-01-01", "2003-01-01", freq="D")) check( assert_type( - pd.IntervalIndex.from_breaks( - pd.Series(pd.date_range("2000-01-01", "2003-01-01", freq="D")) - ), + pd.IntervalIndex.from_breaks(timestamp_series), "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, @@ -426,7 +432,7 @@ def test_interval_index_arrays(): ] ), ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, ) @@ -436,16 +442,16 @@ def test_interval_index_arrays(): pd.IntervalIndex.from_arrays( pd.Series([1, 2, 3, 4]), pd.Series([2, 3, 4, 5]) ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, ) + 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, ) @@ -458,7 +464,8 @@ def test_interval_index_arrays(): pd.Timestamp(2001, 1, 1), pd.Timestamp(2002, 1, 1), pd.Timestamp(2003, 1, 1), - ] + ], + dtype=pd.Timestamp, ), pd.Series( [ @@ -466,10 +473,11 @@ def test_interval_index_arrays(): pd.Timestamp(2002, 1, 1), pd.Timestamp(2003, 1, 1), pd.Timestamp(2004, 1, 1), - ] + ], + dtype=pd.Timestamp, ), ), - pd.IntervalIndex, + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, ) 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 994174ff4..2c6c1a84a 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -7,6 +7,7 @@ ) import numpy as np +import numpy.typing as npt import pandas as pd from typing_extensions import ( TypeAlias, @@ -31,8 +32,6 @@ from pandas._typing import np_ndarray_bool else: - import numpy.typing as npt - np_ndarray_bool = npt.NDArray[np.bool_] PeriodSeries: TypeAlias = pd.Series TimedeltaSeries: TypeAlias = pd.Series @@ -262,56 +261,72 @@ def test_interval() -> None: check(assert_type(hash(interval_td), int), int) interval_index_int = pd.IntervalIndex([interval_i]) - interval_series_int = pd.Series(interval_index_int) - - check(interval_series_int >= interval_i, pd.Series, bool) - check(interval_series_int < interval_i, pd.Series, bool) - check(interval_series_int <= interval_i, pd.Series, bool) - check(interval_series_int > interval_i, pd.Series, bool) - - check(interval_i >= interval_series_int, pd.Series, bool) - check(interval_i < interval_series_int, pd.Series, bool) - check(interval_i <= interval_series_int, pd.Series, bool) - check(interval_i > interval_series_int, pd.Series, bool) - - check(interval_series_int == interval_i, pd.Series, bool) - check(interval_series_int != interval_i, pd.Series, bool) - check( - interval_i - == interval_series_int, # pyright: ignore[reportUnnecessaryComparison] - pd.Series, - bool, + assert_type(interval_index_int >= interval_i, np_ndarray_bool), + np.ndarray, + np.bool_, ) check( - interval_i - != interval_series_int, # pyright: ignore[reportUnnecessaryComparison] - pd.Series, - bool, + 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(interval_index_int >= interval_i, np.ndarray, np.bool_) - check(interval_index_int < interval_i, np.ndarray, np.bool_) - check(interval_index_int <= interval_i, np.ndarray, np.bool_) - check(interval_index_int > interval_i, np.ndarray, np.bool_) - - check(interval_i >= interval_index_int, np.ndarray, np.bool_) - check(interval_i < interval_index_int, np.ndarray, np.bool_) - check(interval_i <= interval_index_int, np.ndarray, np.bool_) - check(interval_i > interval_index_int, 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(interval_index_int == interval_i, np.ndarray, np.bool_) - check(interval_index_int != interval_i, 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( - interval_i - == interval_index_int, # pyright: ignore[reportUnnecessaryComparison] + assert_type( + interval_i == interval_index_int, + np_ndarray_bool, + ), np.ndarray, np.bool_, ) check( - interval_i - != interval_index_int, # pyright: ignore[reportUnnecessaryComparison] + assert_type( + interval_i != interval_index_int, + np_ndarray_bool, + ), np.ndarray, np.bool_, ) From c9ee5801238c011f42a1ce16266e74448fe3e382 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Mon, 7 Nov 2022 15:59:15 +0000 Subject: [PATCH 05/10] ENH: Improve interval --- pandas-stubs/core/indexes/interval.pyi | 54 ++++++++++++++++++++++---- tests/test_indexes.py | 26 +++++++++++-- 2 files changed, 70 insertions(+), 10 deletions(-) diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index 1e71d5dc8..c6f737157 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -165,18 +165,56 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]): copy: bool = ..., dtype: IntervalDtype | None = ..., ) -> IntervalIndex[Interval[pd.Timedelta]]: ... + @overload + @classmethod + def from_tuples( + cls, + data: Sequence[tuple[int, int]] | npt.NDArray[np.integer], + closed: IntervalClosedType = ..., + name: Hashable = ..., + copy: bool = ..., + dtype: IntervalDtype | None = ..., + ) -> IntervalIndex[pd.Interval[int]]: ... + @overload + @classmethod + def from_tuples( + cls, + data: Sequence[tuple[float, float]] | npt.NDArray[np.floating], + 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] + ] + | npt.NDArray[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[int, int]] - | Sequence[tuple[float, float]] - | Sequence[tuple[DatetimeLike, DatetimeLike]] - | npt.NDArray, + data: Sequence[ + tuple[pd.Timedelta, pd.Timedelta] + | tuple[dt.datetime, dt.datetime] + | tuple[np.datetime64, np.datetime64] + ] + | npt.NDArray[np.timedelta64], closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., dtype: IntervalDtype | None = ..., - ) -> IntervalIndex: ... + ) -> IntervalIndex[pd.Interval[pd.Timedelta]]: ... def __contains__(self, key: Any) -> bool: ... def astype(self, dtype: DtypeArg, copy: bool = ...) -> IntervalIndex: ... @property @@ -265,8 +303,10 @@ def interval_range( # type: ignore[misc] name: Hashable = ..., closed: IntervalClosedType = ..., ) -> IntervalIndex[Interval[int]]: ... + +# Overlaps since int is a subclass of float @overload -def interval_range( +def interval_range( # pyright: reportOverlappingOverload=false start: int, *, end: None = ..., @@ -276,7 +316,7 @@ def interval_range( closed: IntervalClosedType = ..., ) -> IntervalIndex[Interval[int]]: ... @overload -def interval_range( +def interval_range( # pyright: reportOverlappingOverload=false *, start: None = ..., end: int, diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 260aedb8b..5afecd449 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -504,13 +504,33 @@ def test_interval_index_arrays(): def test_interval_index_tuples(): + int_arr: npt.NDArray[np.integer] = np.array([[1, 2], [3, 4]], dtype=np.int64) + float_arr: npt.NDArray[np.floating] = np.array([[1, 2], [3, 4]], dtype=np.float_) + 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, ) check( assert_type( - pd.IntervalIndex.from_tuples([(1.0, 2.0), (2.0, 3.0)]), pd.IntervalIndex + pd.IntervalIndex.from_tuples(int_arr), "pd.IntervalIndex[pd.Interval[int]]" + ), + pd.IntervalIndex, + ) + check( + assert_type( + pd.IntervalIndex.from_tuples([(1.0, 2.0), (2.0, 3.0)]), + "pd.IntervalIndex[pd.Interval[float]]", + ), + pd.IntervalIndex, + ) + check( + assert_type( + pd.IntervalIndex.from_tuples(float_arr), + "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, ) @@ -522,7 +542,7 @@ 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, ) From cc45bf1cf0001ee9800a7ee0cea573fb032ff67c Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Mon, 7 Nov 2022 19:01:18 +0000 Subject: [PATCH 06/10] ENH: Improve from tuples --- pandas-stubs/core/indexes/interval.pyi | 17 +++--- tests/test_indexes.py | 84 +++++++++++++++++--------- 2 files changed, 64 insertions(+), 37 deletions(-) diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index c6f737157..ff7b9ef90 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -167,19 +167,20 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]): ) -> IntervalIndex[Interval[pd.Timedelta]]: ... @overload @classmethod - def from_tuples( + def from_tuples( # type:ignore[misc] cls, - data: Sequence[tuple[int, int]] | npt.NDArray[np.integer], + data: Sequence[tuple[int, int]], closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., dtype: IntervalDtype | None = ..., ) -> 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]] | npt.NDArray[np.floating], + data: Sequence[tuple[float, float]], closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., @@ -193,8 +194,7 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]): tuple[pd.Timestamp, pd.Timestamp] | tuple[dt.datetime, dt.datetime] | tuple[np.datetime64, np.datetime64] - ] - | npt.NDArray[np.datetime64], + ], closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., @@ -206,10 +206,9 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]): cls, data: Sequence[ tuple[pd.Timedelta, pd.Timedelta] - | tuple[dt.datetime, dt.datetime] - | tuple[np.datetime64, np.datetime64] - ] - | npt.NDArray[np.timedelta64], + | tuple[dt.timedelta, dt.timedelta] + | tuple[np.timedelta64, np.timedelta64] + ], closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 5afecd449..786c132b6 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -456,24 +456,8 @@ def test_interval_index_arrays(): ), pd.IntervalIndex, ) - left_s_ts = pd.Series( - [ - pd.Timestamp(2000, 1, 1), - pd.Timestamp(2001, 1, 1), - pd.Timestamp(2002, 1, 1), - pd.Timestamp(2003, 1, 1), - ], - dtype=pd.Timestamp, - ) - right_s_ts = pd.Series( - [ - pd.Timestamp(2001, 1, 1), - pd.Timestamp(2002, 1, 1), - pd.Timestamp(2003, 1, 1), - pd.Timestamp(2004, 1, 1), - ], - dtype=pd.Timestamp, - ) + 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(left_s_ts, right_s_ts), @@ -504,9 +488,6 @@ def test_interval_index_arrays(): def test_interval_index_tuples(): - int_arr: npt.NDArray[np.integer] = np.array([[1, 2], [3, 4]], dtype=np.int64) - float_arr: npt.NDArray[np.floating] = np.array([[1, 2], [3, 4]], dtype=np.float_) - check( assert_type( pd.IntervalIndex.from_tuples([(1, 2), (2, 3)]), @@ -516,21 +497,32 @@ def test_interval_index_tuples(): ) check( assert_type( - pd.IntervalIndex.from_tuples(int_arr), "pd.IntervalIndex[pd.Interval[int]]" + pd.IntervalIndex.from_tuples([(1.0, 2.0), (2.0, 3.0)]), + "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, ) check( assert_type( - pd.IntervalIndex.from_tuples([(1.0, 2.0), (2.0, 3.0)]), - "pd.IntervalIndex[pd.Interval[float]]", + pd.IntervalIndex.from_tuples( + [ + (pd.Timestamp(2000, 1, 1), pd.Timestamp(2001, 1, 1)), + (pd.Timestamp(2001, 1, 1), pd.Timestamp(2002, 1, 1)), + ] + ), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, ) check( assert_type( - pd.IntervalIndex.from_tuples(float_arr), - "pd.IntervalIndex[pd.Interval[float]]", + 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, ) @@ -538,11 +530,47 @@ def test_interval_index_tuples(): assert_type( pd.IntervalIndex.from_tuples( [ - (pd.Timestamp(2000, 1, 1), pd.Timestamp(2001, 1, 1)), - (pd.Timestamp(2001, 1, 1), pd.Timestamp(2002, 1, 1)), + (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, ) + 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, + ) + 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, + ) + 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, + ) From 657c1d2e2e0e30d3b912e28be15dafda07d35aa4 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Wed, 23 Nov 2022 18:35:49 +0000 Subject: [PATCH 07/10] CLN: Final clean ups --- pandas-stubs/_libs/interval.pyi | 60 +++----------------------- pandas-stubs/core/indexes/interval.pyi | 2 +- tests/test_scalars.py | 49 ++++++++++++++++++--- 3 files changed, 52 insertions(+), 59 deletions(-) diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index d044697f3..8d7ff4d69 100644 --- a/pandas-stubs/_libs/interval.pyi +++ b/pandas-stubs/_libs/interval.pyi @@ -158,71 +158,31 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def overlaps(self: Interval[float], other: Interval[int]) -> bool: ... @overload - def __gt__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[misc] + def __gt__(self, other: Interval[_OrderableT]) -> bool: ... @overload def __gt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] - # @overload - # def __gt__(self:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... - # @overload - # def __gt__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... - # @overload - # def __gt__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... @overload - def __gt__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[misc] + def __gt__(self, other: Series[_OrderableT]) -> Series[bool]: ... @overload - def __gt__(self, other: object) -> Literal[False]: ... - @overload - def __lt__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[misc] + def __lt__(self, other: Interval[_OrderableT]) -> bool: ... @overload def __lt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] - # @overload - # def __lt__(self:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... - # @overload - # def __lt__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... - # @overload - # def __lt__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... - @overload - def __lt__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[misc] @overload - def __lt__(self, other: object) -> Literal[False]: ... + def __lt__(self, other: Series[_OrderableT]) -> Series[bool]: ... @overload - def __ge__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[misc] + def __ge__(self, other: Interval[_OrderableT]) -> bool: ... @overload def __ge__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] - # @overload - # def __ge__(self:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... - # @overload - # def __ge__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... - # @overload - # def __ge__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... @overload - def __ge__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[misc] + def __ge__(self, other: Series[_OrderableT]) -> Series[bool]: ... @overload - def __ge__(self, other: object) -> Literal[False]: ... - @overload - def __le__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[misc] + def __le__(self, other: Interval[_OrderableT]) -> bool: ... @overload def __le__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] - # @overload - # def __le__(self:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... - # @overload - # def __le__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... - # @overload - # def __le__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... - # @overload - # def __le__(self, other: Series[_OrderableT]) -> Series[bool]: ... - @overload - def __le__(self, other: object) -> Literal[False]: ... @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:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... - # @overload - # def __eq__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... - # @overload - # def __eq__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... @overload def __eq__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[misc] @overload @@ -231,12 +191,6 @@ class Interval(IntervalMixin, Generic[_OrderableT]): 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:Interval[float], other: IntervalIndex[Interval[float]]) -> np_ndarray_bool: ... - # @overload - # def __ne__(self:Interval[Timestamp], other: IntervalIndex[Interval[Timestamp]]) -> np_ndarray_bool: ... - # @overload - # def __ne__(self:Interval[Timedelta], other: IntervalIndex[Interval[Timedelta]]) -> np_ndarray_bool: ... @overload def __ne__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[misc] @overload diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index ff7b9ef90..39759d769 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -167,7 +167,7 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]): ) -> IntervalIndex[Interval[pd.Timedelta]]: ... @overload @classmethod - def from_tuples( # type:ignore[misc] + def from_tuples( cls, data: Sequence[tuple[int, int]], closed: IntervalClosedType = ..., diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 6c67c9ac0..53dbd28ba 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -152,6 +152,17 @@ def test_interval() -> None: 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( @@ -216,6 +227,7 @@ def test_interval() -> None: 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( @@ -240,27 +252,59 @@ def test_interval() -> None: 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) @@ -268,11 +312,6 @@ def test_interval() -> None: 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) - check(assert_type(hash(interval_i), int), int) - check(assert_type(hash(interval_f), int), int) - check(assert_type(hash(interval_ts), int), int) - check(assert_type(hash(interval_td), int), int) - interval_index_int = pd.IntervalIndex([interval_i]) check( assert_type(interval_index_int >= interval_i, np_ndarray_bool), From 684ef2eacb4c978b8a0f833e478463fe04e3d43e Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Thu, 24 Nov 2022 07:08:47 +0000 Subject: [PATCH 08/10] CLN: Final clean ups --- pandas-stubs/core/indexes/interval.pyi | 26 ++++++++--- tests/test_indexes.py | 62 ++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 7 deletions(-) diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index 39759d769..9792522f0 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -17,7 +17,10 @@ from pandas.core.series import ( TimedeltaSeries, TimestampSeries, ) -from typing_extensions import TypeAlias +from typing_extensions import ( + Never, + TypeAlias, +) from pandas._libs.interval import ( Interval as Interval, @@ -259,31 +262,31 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]): @overload def __gt__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ... # type: ignore[misc] @overload - def __gt__(self, other: object) -> Literal[False]: ... + 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) -> Literal[False]: ... + 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) -> Literal[False]: ... + 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) -> Literal[False]: ... + 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]: ... + def __eq__(self, other: object) -> Never: ... @overload # type: ignore[override] def __ne__(self, other: IntervalT | IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] @overload @@ -383,6 +386,15 @@ def interval_range( 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 = ..., @@ -391,7 +403,7 @@ def interval_range( freq: str | BaseOffset | None = ..., name: Hashable = ..., closed: IntervalClosedType = ..., -) -> IntervalIndex[Interval[pd.Timestamp]]: ... +) -> IntervalIndex[Interval[pd.Timedelta]]: ... @overload def interval_range( start: _TimedeltaLike, diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 786c132b6..cd4d5c9ca 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -194,6 +194,7 @@ def test_interval_range(): check( assert_type(pd.interval_range(0, 10), "pd.IntervalIndex[pd.Interval[int]]"), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -201,10 +202,12 @@ def test_interval_range(): "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( @@ -212,6 +215,7 @@ def test_interval_range(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -221,6 +225,7 @@ def test_interval_range(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -228,6 +233,7 @@ def test_interval_range(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -237,6 +243,7 @@ def test_interval_range(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -248,6 +255,7 @@ def test_interval_range(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -255,6 +263,32 @@ def test_interval_range(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, + ) + + check( + 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.interval_range(start=pd.Timedelta("1D"), periods=10, freq="D"), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", + ), + pd.IntervalIndex, + pd.Interval, ) @@ -265,6 +299,7 @@ def test_interval_index_breaks(): "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -272,6 +307,7 @@ def test_interval_index_breaks(): "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -281,6 +317,7 @@ def test_interval_index_breaks(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -288,6 +325,7 @@ def test_interval_index_breaks(): "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", ), pd.IntervalIndex, + pd.Interval, ) check( @@ -296,6 +334,7 @@ def test_interval_index_breaks(): "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -305,6 +344,7 @@ def test_interval_index_breaks(): "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, ) np_ndarray_dt64 = np.array( [ @@ -321,6 +361,7 @@ def test_interval_index_breaks(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) np_ndarray_td64 = np.array( [ @@ -337,6 +378,7 @@ def test_interval_index_breaks(): "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -352,6 +394,7 @@ def test_interval_index_breaks(): "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( @@ -360,6 +403,7 @@ def test_interval_index_breaks(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -374,6 +418,7 @@ def test_interval_index_breaks(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) @@ -384,6 +429,7 @@ def test_interval_index_arrays(): "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -391,6 +437,7 @@ def test_interval_index_arrays(): "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -401,6 +448,7 @@ def test_interval_index_arrays(): "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -411,6 +459,7 @@ def test_interval_index_arrays(): "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, ) left_dt64_arr: npt.NDArray[np.datetime64] = np.array( [ @@ -436,6 +485,7 @@ def test_interval_index_arrays(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( @@ -446,6 +496,7 @@ def test_interval_index_arrays(): "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) @@ -455,6 +506,7 @@ def test_interval_index_arrays(): "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")) @@ -464,6 +516,7 @@ def test_interval_index_arrays(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -484,6 +537,7 @@ def test_interval_index_arrays(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) @@ -494,6 +548,7 @@ def test_interval_index_tuples(): "pd.IntervalIndex[pd.Interval[int]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -501,6 +556,7 @@ def test_interval_index_tuples(): "pd.IntervalIndex[pd.Interval[float]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -513,6 +569,7 @@ def test_interval_index_tuples(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -525,6 +582,7 @@ def test_interval_index_tuples(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -537,6 +595,7 @@ def test_interval_index_tuples(): "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -549,6 +608,7 @@ def test_interval_index_tuples(): "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -561,6 +621,7 @@ def test_interval_index_tuples(): "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", ), pd.IntervalIndex, + pd.Interval, ) check( assert_type( @@ -573,4 +634,5 @@ def test_interval_index_tuples(): "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", ), pd.IntervalIndex, + pd.Interval, ) From f01a3882f2888cfb00e67e8bb3c331f72edc319a Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Thu, 24 Nov 2022 07:30:22 +0000 Subject: [PATCH 09/10] CLN: Final fixes --- pandas-stubs/_libs/interval.pyi | 8 ++++---- pandas-stubs/core/indexes/interval.pyi | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pandas-stubs/_libs/interval.pyi b/pandas-stubs/_libs/interval.pyi index 8d7ff4d69..69e79b7c8 100644 --- a/pandas-stubs/_libs/interval.pyi +++ b/pandas-stubs/_libs/interval.pyi @@ -160,25 +160,25 @@ class Interval(IntervalMixin, Generic[_OrderableT]): @overload def __gt__(self, other: Interval[_OrderableT]) -> bool: ... @overload - def __gt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc] + 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: ... # type: ignore[misc] + 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: ... # type: ignore[misc] + 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: ... # type: ignore[misc] + def __le__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... @overload def __eq__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[misc] @overload diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index 9792522f0..8234f7a07 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -281,7 +281,7 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]): def __lt__(self, other: pd.Series[IntervalT]) -> bool: ... # type: ignore[misc] @overload def __lt__(self, other: object) -> Never: ... - @overload # type: ignore[override] + @overload 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] From ee0003ba365ff42589bee750bbaa7e0cef9feb89 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Thu, 24 Nov 2022 18:25:02 +0000 Subject: [PATCH 10/10] CLN: Final fixes --- .flake8 | 3 --- pandas-stubs/core/indexes/interval.pyi | 10 ++++++---- 2 files changed, 6 insertions(+), 7 deletions(-) delete mode 100644 .flake8 diff --git a/.flake8 b/.flake8 deleted file mode 100644 index fe708d67b..000000000 --- a/.flake8 +++ /dev/null @@ -1,3 +0,0 @@ -[flake8] -per-file-ignores = - tests/test_scalars.py: W503 \ No newline at end of file diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index 8234f7a07..5842ace9e 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -1,6 +1,5 @@ import datetime as dt from typing import ( - Any, Generic, Hashable, Literal, @@ -217,7 +216,10 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]): copy: bool = ..., dtype: IntervalDtype | None = ..., ) -> IntervalIndex[pd.Interval[pd.Timedelta]]: ... - def __contains__(self, key: Any) -> bool: ... + @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: ... @@ -281,12 +283,12 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]): def __lt__(self, other: pd.Series[IntervalT]) -> bool: ... # type: ignore[misc] @overload def __lt__(self, other: object) -> Never: ... - @overload + @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) -> Never: ... + 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