From becf720692f4ef64c0665f4f791e9bfe3676dce5 Mon Sep 17 00:00:00 2001 From: Amotz Date: Sat, 30 Jul 2022 20:20:38 +0300 Subject: [PATCH 1/2] Fixed typing on IntervalIndex functions --- pandas-stubs/core/indexes/interval.pyi | 62 ++++++++++++++------------ tests/test_interval_index.py | 23 ++++++++++ 2 files changed, 56 insertions(+), 29 deletions(-) create mode 100644 tests/test_interval_index.py diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index ab052fe30..d821aada3 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -1,5 +1,9 @@ +from typing import ( + Hashable, + Literal, +) + import numpy as np -from pandas.core.indexes.base import Index from pandas.core.indexes.extension import ExtensionIndex from pandas._libs.interval import ( @@ -8,6 +12,7 @@ from pandas._libs.interval import ( ) from pandas._typing import DtypeArg +from pandas.core.dtypes.dtypes import IntervalDtype as IntervalDtype from pandas.core.dtypes.generic import ABCSeries as ABCSeries class SetopCheck: @@ -19,38 +24,46 @@ class IntervalIndex(IntervalMixin, ExtensionIndex): def __new__( cls, data, - closed=..., - dtype=..., + closed: Literal["left", "right", "both", "neither"] = ..., + dtype: IntervalDtype | None = ..., copy: bool = ..., - name=..., + name: Hashable = ..., verify_integrity: bool = ..., ): ... @classmethod def from_breaks( - cls, breaks, closed: str = ..., name=..., copy: bool = ..., dtype=... - ): ... + cls, + breaks, + closed: Literal["left", "right", "both", "neither"] = ..., + name: Hashable = ..., + copy: bool = ..., + dtype: IntervalDtype | None = ..., + ) -> IntervalIndex: ... @classmethod def from_arrays( - cls, left, right, closed: str = ..., name=..., copy: bool = ..., dtype=... - ): ... + cls, + left, + right, + closed: Literal["left", "right", "both", "neither"] = ..., + name: Hashable = ..., + copy: bool = ..., + dtype: IntervalDtype | None = ..., + ) -> IntervalIndex: ... @classmethod def from_tuples( - cls, data, closed: str = ..., name=..., copy: bool = ..., dtype=... - ): ... - def __contains__(self, key) -> bool: ... - def values(self): ... - def __array_wrap__(self, result, context=...): ... - def __reduce__(self): ... - def astype(self, dtype: DtypeArg, copy: bool = ...) -> Index: ... + cls, + data, + closed: Literal["left", "right", "both", "neither"] = ..., + name: Hashable = ..., + copy: bool = ..., + dtype: IntervalDtype | None = ..., + ) -> IntervalIndex: ... + def astype(self, dtype: DtypeArg, copy: bool = ...) -> IntervalIndex: ... @property def inferred_type(self) -> str: ... def memory_usage(self, deep: bool = ...) -> int: ... - def is_monotonic(self) -> bool: ... - def is_monotonic_increasing(self) -> bool: ... - def is_monotonic_decreasing(self) -> bool: ... - def is_unique(self): ... @property - def is_overlapping(self): ... + def is_overlapping(self) -> bool: ... def get_loc( self, key, method: str | None = ..., tolerance=... ) -> int | slice | np.ndarray: ... @@ -65,15 +78,6 @@ class IntervalIndex(IntervalMixin, ExtensionIndex): self, targetArrayLike ) -> tuple[np.ndarray, np.ndarray]: ... def get_value(self, series: ABCSeries, key): ... - def where(self, cond, other=...): ... - def delete(self, loc): ... - def insert(self, loc, item): ... - def take( - self, indices, axis: int = ..., allow_fill: bool = ..., fill_value=..., **kwargs - ): ... - def __getitem__(self, value): ... - def argsort(self, *args, **kwargs): ... - def equals(self, other) -> bool: ... @property def is_all_dates(self) -> bool: ... def __lt__(self, other): ... diff --git a/tests/test_interval_index.py b/tests/test_interval_index.py new file mode 100644 index 000000000..5d8c5f0a8 --- /dev/null +++ b/tests/test_interval_index.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +import pandas as pd +from typing_extensions import assert_type + +from tests import check + + +def test_from_breaks() -> None: + ind1: pd.IntervalIndex = pd.IntervalIndex.from_breaks([0, 1, 2, 3], name="test") + ind2: pd.IntervalIndex = pd.IntervalIndex.from_breaks([0, 1, 2, 3], closed="right", name=123) + +def test_from_arrays() -> None: + ind1: pd.IntervalIndex = pd.IntervalIndex.from_arrays([0, 1, 2], [1, 2, 3], name="test") + ind2: pd.IntervalIndex = pd.IntervalIndex.from_arrays([0, 1, 2], [1, 2, 3], closed="right", name=123) + +def test_from_tuples() -> None: + ind1: pd.IntervalIndex = pd.IntervalIndex.from_tuples([(0, 1), (1, 2), (2, 3)], name="test") + ind2: pd.IntervalIndex = pd.IntervalIndex.from_tuples([(0, 1), (1, 2), (2, 3)], closed="right", name=123) + +def test_is_overlapping() -> None: + ind = pd.IntervalIndex.from_tuples([(0, 2), (1, 3), (4, 5)]) + check(assert_type(ind.is_overlapping, bool), bool) From 27b2dee521d490c3276f036553ff59c4e18385b8 Mon Sep 17 00:00:00 2001 From: Amotz Date: Sun, 31 Jul 2022 18:59:30 +0300 Subject: [PATCH 2/2] Used IntervalClosedType and fixes tests style issues --- pandas-stubs/core/indexes/interval.pyi | 18 +++++++++--------- tests/test_interval_index.py | 23 ++++++++++++++++++----- 2 files changed, 27 insertions(+), 14 deletions(-) diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index d821aada3..974720a98 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -1,7 +1,4 @@ -from typing import ( - Hashable, - Literal, -) +from typing import Hashable import numpy as np from pandas.core.indexes.extension import ExtensionIndex @@ -10,7 +7,10 @@ from pandas._libs.interval import ( Interval as Interval, IntervalMixin as IntervalMixin, ) -from pandas._typing import DtypeArg +from pandas._typing import ( + DtypeArg, + IntervalClosedType, +) from pandas.core.dtypes.dtypes import IntervalDtype as IntervalDtype from pandas.core.dtypes.generic import ABCSeries as ABCSeries @@ -24,7 +24,7 @@ class IntervalIndex(IntervalMixin, ExtensionIndex): def __new__( cls, data, - closed: Literal["left", "right", "both", "neither"] = ..., + closed: IntervalClosedType = ..., dtype: IntervalDtype | None = ..., copy: bool = ..., name: Hashable = ..., @@ -34,7 +34,7 @@ class IntervalIndex(IntervalMixin, ExtensionIndex): def from_breaks( cls, breaks, - closed: Literal["left", "right", "both", "neither"] = ..., + closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., dtype: IntervalDtype | None = ..., @@ -44,7 +44,7 @@ class IntervalIndex(IntervalMixin, ExtensionIndex): cls, left, right, - closed: Literal["left", "right", "both", "neither"] = ..., + closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., dtype: IntervalDtype | None = ..., @@ -53,7 +53,7 @@ class IntervalIndex(IntervalMixin, ExtensionIndex): def from_tuples( cls, data, - closed: Literal["left", "right", "both", "neither"] = ..., + closed: IntervalClosedType = ..., name: Hashable = ..., copy: bool = ..., dtype: IntervalDtype | None = ..., diff --git a/tests/test_interval_index.py b/tests/test_interval_index.py index 5d8c5f0a8..ec26019b6 100644 --- a/tests/test_interval_index.py +++ b/tests/test_interval_index.py @@ -8,15 +8,28 @@ def test_from_breaks() -> None: ind1: pd.IntervalIndex = pd.IntervalIndex.from_breaks([0, 1, 2, 3], name="test") - ind2: pd.IntervalIndex = pd.IntervalIndex.from_breaks([0, 1, 2, 3], closed="right", name=123) + ind2: pd.IntervalIndex = pd.IntervalIndex.from_breaks( + [0, 1, 2, 3], closed="right", name=123 + ) + def test_from_arrays() -> None: - ind1: pd.IntervalIndex = pd.IntervalIndex.from_arrays([0, 1, 2], [1, 2, 3], name="test") - ind2: pd.IntervalIndex = pd.IntervalIndex.from_arrays([0, 1, 2], [1, 2, 3], closed="right", name=123) + ind1: pd.IntervalIndex = pd.IntervalIndex.from_arrays( + [0, 1, 2], [1, 2, 3], name="test" + ) + ind2: pd.IntervalIndex = pd.IntervalIndex.from_arrays( + [0, 1, 2], [1, 2, 3], closed="right", name=123 + ) + def test_from_tuples() -> None: - ind1: pd.IntervalIndex = pd.IntervalIndex.from_tuples([(0, 1), (1, 2), (2, 3)], name="test") - ind2: pd.IntervalIndex = pd.IntervalIndex.from_tuples([(0, 1), (1, 2), (2, 3)], closed="right", name=123) + ind1: pd.IntervalIndex = pd.IntervalIndex.from_tuples( + [(0, 1), (1, 2), (2, 3)], name="test" + ) + ind2: pd.IntervalIndex = pd.IntervalIndex.from_tuples( + [(0, 1), (1, 2), (2, 3)], closed="right", name=123 + ) + def test_is_overlapping() -> None: ind = pd.IntervalIndex.from_tuples([(0, 2), (1, 3), (4, 5)])