Fix __getitem__() for various index types (pandas-dev#537)

* Fix __getitem__() for various index types

* create TypeVar for the mixin

Author: Dr-Irv

pandas-stubs/core/algorithms.pyi

@@ -25,7 +25,7 @@ def unique(values: PeriodIndex) -> PeriodIndex: ...  # type: ignore[misc] # pyri
 @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]: ...
 @overload
 def unique(values: Index) -> np.ndarray: ...
 @overload

pandas-stubs/core/indexes/base.pyi

@@ -7,7 +7,9 @@ from collections.abc import (
 )
 from typing import (
     ClassVar,
+    Generic,
     Literal,
+    TypeVar,
     overload,
 )
 
@@ -27,6 +29,7 @@ from pandas.core.strings import StringMethods
 from typing_extensions import Never
 
 from pandas._typing import (
+    S1,
     T1,
     Dtype,
     DtypeArg,
@@ -48,6 +51,23 @@ class InvalidIndexError(Exception): ...
 
 _str = str
 
+_IndexGetitemMixinT = TypeVar("_IndexGetitemMixinT", bound=_IndexGetitemMixin)
+
+class _IndexGetitemMixin(Generic[S1]):
+    @overload
+    def __getitem__(
+        self: _IndexGetitemMixinT,
+        idx: slice
+        | np_ndarray_anyint
+        | Sequence[int]
+        | Index
+        | Series[bool]
+        | Sequence[bool]
+        | np_ndarray_bool,
+    ) -> _IndexGetitemMixinT: ...
+    @overload
+    def __getitem__(self, idx: int) -> S1: ...
+
 class Index(IndexOpsMixin, PandasObject):
     __hash__: ClassVar[None]  # type: ignore[assignment]
     @overload
@@ -158,7 +178,7 @@ class Index(IndexOpsMixin, PandasObject):
     __bool__ = ...
     def union(self, other: list[HashableT] | Index, sort=...) -> Index: ...
     def intersection(self, other: list[T1] | Index, sort: bool = ...) -> Index: ...
-    def difference(self, other: list | Index) -> Index: ...
+    def difference(self, other: list | Index, sort: bool | None = None) -> Index: ...
     def symmetric_difference(
         self, other: list[T1] | Index, result_name=..., sort=...
     ) -> Index: ...
@@ -191,7 +211,13 @@ class Index(IndexOpsMixin, PandasObject):
     @overload
     def __getitem__(
         self: IndexT,
-        idx: slice | np_ndarray_anyint | Index | Series[bool] | np_ndarray_bool,
+        idx: slice
+        | np_ndarray_anyint
+        | Sequence[int]
+        | Index
+        | Series[bool]
+        | Sequence[bool]
+        | np_ndarray_bool,
     ) -> IndexT: ...
     @overload
     def __getitem__(self, idx: int | tuple[np_ndarray_anyint, ...]) -> Scalar: ...

pandas-stubs/core/indexes/datetimelike.pyi

@@ -1,5 +1,5 @@
 from pandas.core.indexes.extension import ExtensionIndex
-from pandas.core.indexes.numeric import Int64Index
+from pandas.core.indexes.timedeltas import TimedeltaIndex
 
 from pandas._libs.tslibs import BaseOffset
 
@@ -10,28 +10,10 @@ class DatetimeIndexOpsMixin(ExtensionIndex):
     def freqstr(self) -> str | None: ...
     @property
     def is_all_dates(self) -> bool: ...
-    @property
-    def values(self): ...
-    def __array_wrap__(self, result, context=...): ...
-    def equals(self, other) -> bool: ...
-    def __contains__(self, key): ...
-    def sort_values(self, return_indexer: bool = ..., ascending: bool = ...): ...
-    def take(
-        self, indices, axis: int = ..., allow_fill: bool = ..., fill_value=..., **kwargs
-    ): ...
-    def tolist(self) -> list: ...
     def min(self, axis=..., skipna: bool = ..., *args, **kwargs): ...
     def argmin(self, axis=..., skipna: bool = ..., *args, **kwargs): ...
     def max(self, axis=..., skipna: bool = ..., *args, **kwargs): ...
     def argmax(self, axis=..., skipna: bool = ..., *args, **kwargs): ...
-    def isin(self, values, level=...): ...
-    def where(self, cond, other=...): ...
-    def shift(self, periods: int = ..., freq=...): ...
-    def delete(self, loc): ...
+    def __rsub__(self, other: DatetimeIndexOpsMixin) -> TimedeltaIndex: ...
 
-class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin, Int64Index):
-    def difference(self, other, sort=...): ...
-    def intersection(self, other, sort: bool = ...): ...
-    def join(
-        self, other, *, how: str = ..., level=..., return_indexers=..., sort=...
-    ): ...
+class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin): ...

pandas-stubs/core/indexes/datetimes.pyi

@@ -17,6 +17,7 @@ from pandas import (
 )
 from pandas.core.indexes.accessors import DatetimeIndexProperties
 from pandas.core.indexes.api import Float64Index
+from pandas.core.indexes.base import _IndexGetitemMixin
 from pandas.core.indexes.datetimelike import DatetimeTimedeltaMixin
 from pandas.core.series import (
     TimedeltaSeries,
@@ -34,7 +35,12 @@ from pandas.core.dtypes.dtypes import DatetimeTZDtype
 
 from pandas.tseries.offsets import BaseOffset
 
-class DatetimeIndex(DatetimeTimedeltaMixin, DatetimeIndexProperties):
+# type ignore needed because of __getitem__()
+class DatetimeIndex(  # type: ignore[misc]
+    _IndexGetitemMixin[Timestamp],
+    DatetimeTimedeltaMixin,
+    DatetimeIndexProperties,
+):
     def __init__(
         self,
         data: ArrayLike | AnyArrayLike | list | tuple,
@@ -53,11 +59,11 @@ class DatetimeIndex(DatetimeTimedeltaMixin, DatetimeIndexProperties):
     def __reduce__(self): ...
     # various ignores needed for mypy, as we do want to restrict what can be used in
     # arithmetic for these types
-    @overload  # type: ignore[override]
+    @overload
     def __add__(self, other: TimedeltaSeries) -> TimestampSeries: ...
     @overload
     def __add__(self, other: Timedelta | TimedeltaIndex) -> DatetimeIndex: ...
-    @overload  # type: ignore[override]
+    @overload
     def __sub__(self, other: TimedeltaSeries) -> TimestampSeries: ...
     @overload
     def __sub__(self, other: Timedelta | TimedeltaIndex) -> DatetimeIndex: ...

pandas-stubs/core/indexes/extension.pyi

@@ -1,14 +1,3 @@
-from typing import Literal
-
 from pandas.core.indexes.base import Index
 
-class ExtensionIndex(Index):
-    def __iter__(self): ...
-    def dropna(self, how: Literal["any", "all"] = ...): ...
-    def repeat(self, repeats, axis=...): ...
-    def take(
-        self, indices, axis: int = ..., allow_fill: bool = ..., fill_value=..., **kwargs
-    ): ...
-    def unique(self, level=...): ...
-    def map(self, mapper, na_action=...): ...
-    def astype(self, dtype, copy: bool = ...): ...
+class ExtensionIndex(Index): ...

pandas-stubs/core/indexes/interval.pyi

@@ -13,8 +13,8 @@ from typing import (
 import numpy as np
 import pandas as pd
 from pandas import Index
-from pandas.core.indexes.extension import ExtensionIndex
 from pandas.core.series import (
+    Series,
     TimedeltaSeries,
     TimestampSeries,
 )
@@ -32,6 +32,7 @@ from pandas._typing import (
     IntervalClosedType,
     IntervalT,
     Label,
+    np_ndarray_anyint,
     np_ndarray_bool,
     npt,
 )
@@ -68,7 +69,7 @@ _EdgesTimedelta: TypeAlias = Union[
 _TimestampLike: TypeAlias = Union[pd.Timestamp, np.datetime64, dt.datetime]
 _TimedeltaLike: TypeAlias = Union[pd.Timedelta, np.timedelta64, dt.timedelta]
 
-class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]):
+class IntervalIndex(IntervalMixin, Generic[IntervalT]):
     def __new__(
         cls,
         data: Sequence[IntervalT],
@@ -78,10 +79,9 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]):
         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]
+    def from_breaks(
         cls,
         breaks: _EdgesInt,
         closed: IntervalClosedType = ...,
@@ -119,10 +119,9 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]):
         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]
+    def from_arrays(
         cls,
         left: _EdgesInt,
         right: _EdgesInt,
@@ -250,41 +249,51 @@ class IntervalIndex(IntervalMixin, ExtensionIndex, Generic[IntervalT]):
     @property
     def length(self) -> Index: ...
     def get_value(self, series: ABCSeries, key): ...
+    @overload
+    def __getitem__(
+        self,
+        idx: slice
+        | np_ndarray_anyint
+        | Sequence[int]
+        | Index
+        | Series[bool]
+        | np_ndarray_bool,
+    ) -> IntervalIndex[IntervalT]: ...
+    @overload
+    def __getitem__(self, idx: int) -> IntervalT: ...
     @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]
+    @overload
     def __gt__(
         self, other: IntervalT | IntervalIndex[IntervalT]
     ) -> np_ndarray_bool: ...
     @overload
     def __gt__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ...
-    @overload  # type: ignore[override]
+    @overload
     def __ge__(
         self, other: IntervalT | IntervalIndex[IntervalT]
     ) -> np_ndarray_bool: ...
     @overload
     def __ge__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ...
-    @overload  # type: ignore[override]
+    @overload
     def __le__(
         self, other: IntervalT | IntervalIndex[IntervalT]
     ) -> np_ndarray_bool: ...
     @overload
     def __le__(self, other: pd.Series[IntervalT]) -> pd.Series[bool]: ...
-    @overload  # type: ignore[override]
+    @overload
     def __lt__(
         self, other: IntervalT | IntervalIndex[IntervalT]
     ) -> np_ndarray_bool: ...
     @overload
-    def __lt__(self, other: pd.Series[IntervalT]) -> bool: ...  # type: ignore[misc]
-    @overload  # type: ignore[override]
+    def __lt__(self, other: pd.Series[IntervalT]) -> bool: ...
+    @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]
     @overload
     def __eq__(self, other: object) -> Literal[False]: ...
-    @overload  # type: ignore[override]
+    @overload
     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]

pandas-stubs/core/indexes/multi.pyi

@@ -3,7 +3,10 @@ from collections.abc import (
     Hashable,
     Sequence,
 )
-from typing import Literal
+from typing import (
+    Literal,
+    overload,
+)
 
 import numpy as np
 import pandas as pd
@@ -13,6 +16,7 @@ from pandas._typing import (
     T1,
     DtypeArg,
     HashableT,
+    np_ndarray_anyint,
     np_ndarray_bool,
 )
 
@@ -105,7 +109,19 @@ class MultiIndex(Index):
     @property
     def levshape(self): ...
     def __reduce__(self): ...
-    def __getitem__(self, key): ...
+    @overload  # type: ignore[override]
+    def __getitem__(
+        self,
+        idx: slice
+        | np_ndarray_anyint
+        | Sequence[int]
+        | Index
+        | pd.Series[bool]
+        | Sequence[bool]
+        | np_ndarray_bool,
+    ) -> MultiIndex: ...
+    @overload
+    def __getitem__(self, key: int) -> tuple: ...
     def take(
         self, indices, axis: int = ..., allow_fill: bool = ..., fill_value=..., **kwargs
     ): ...

pandas-stubs/core/indexes/period.pyi

@@ -5,18 +5,25 @@ import numpy as np
 import pandas as pd
 from pandas import Index
 from pandas.core.indexes.accessors import PeriodIndexFieldOps
+from pandas.core.indexes.base import _IndexGetitemMixin
 from pandas.core.indexes.datetimelike import (
     DatetimeIndexOpsMixin as DatetimeIndexOpsMixin,
 )
-from pandas.core.indexes.numeric import Int64Index
-from pandas.core.series import OffsetSeries
+from pandas.core.indexes.timedeltas import TimedeltaIndex
 
 from pandas._libs.tslibs import (
     BaseOffset,
+    NaTType,
     Period,
 )
+from pandas._libs.tslibs.period import _PeriodAddSub
 
-class PeriodIndex(DatetimeIndexOpsMixin, Int64Index, PeriodIndexFieldOps):
+# type ignore needed because of __getitem__()
+class PeriodIndex(  # type: ignore[misc]
+    _IndexGetitemMixin[Period],
+    DatetimeIndexOpsMixin,
+    PeriodIndexFieldOps,
+):
     def __new__(
         cls,
         data=...,
@@ -31,11 +38,22 @@ class PeriodIndex(DatetimeIndexOpsMixin, Int64Index, PeriodIndexFieldOps):
     @property
     def values(self): ...
     def __contains__(self, key) -> bool: ...
-    # Override due to supertype incompatibility which has it for NumericIndex or complex.
-    @overload  # type: ignore[override]
+    @overload
     def __sub__(self, other: Period) -> Index: ...
     @overload
-    def __sub__(self, other: PeriodIndex) -> OffsetSeries: ...
+    def __sub__(self, other: PeriodIndex) -> Index: ...
+    @overload
+    def __sub__(self, other: _PeriodAddSub) -> PeriodIndex: ...
+    @overload
+    def __sub__(self, other: NaTType) -> NaTType: ...
+    @overload
+    def __sub__(self, other: TimedeltaIndex | pd.Timedelta) -> PeriodIndex: ...
+    @overload  # type: ignore[override]
+    def __rsub__(self, other: Period) -> Index: ...
+    @overload
+    def __rsub__(self, other: PeriodIndex) -> Index: ...
+    @overload
+    def __rsub__(self, other: NaTType) -> NaTType: ...
     def __array__(self, dtype=...) -> np.ndarray: ...
     def __array_wrap__(self, result, context=...): ...
     def asof_locs(self, where, mask): ...