TYP: Typing for ExtensionArray.__getitem__

pandas/_typing.py

@@ -211,7 +211,6 @@
 # TODO: add Ellipsis, see
 # https://github.com/python/typing/issues/684#issuecomment-548203158
 # https://bugs.python.org/issue41810
-PositionalIndexer = Union[int, np.integer, slice, Sequence[int], np.ndarray]
-PositionalIndexer2D = Union[
-    PositionalIndexer, Tuple[PositionalIndexer, PositionalIndexer]
-]
+PositionalIndexer = Union[int, np.integer, slice, List[int], np.ndarray]
+PositionalIndexerTuple = Tuple[PositionalIndexer, PositionalIndexer]
+PositionalIndexer2D = Union[PositionalIndexer, PositionalIndexerTuple]

pandas/core/arrays/_mixins.py

@@ -6,6 +6,7 @@
     Sequence,
     TypeVar,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -15,6 +16,7 @@
 from pandas._typing import (
     F,
     PositionalIndexer2D,
+    PositionalIndexerTuple,
     Shape,
     type_t,
 )
@@ -204,6 +206,17 @@ def __setitem__(self, key, value):
     def _validate_setitem_value(self, value):
         return value
 
+    @overload
+    def __getitem__(self, key: int | np.integer) -> Any:
+        ...
+
+    @overload
+    def __getitem__(
+        self: NDArrayBackedExtensionArray,
+        key: slice | np.ndarray | list[int] | PositionalIndexerTuple,
+    ) -> NDArrayBackedExtensionArray:
+        ...
+
     def __getitem__(
         self: NDArrayBackedExtensionArrayT,
         key: PositionalIndexer2D,

pandas/core/arrays/base.py

@@ -17,6 +17,7 @@
     Sequence,
     TypeVar,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -291,8 +292,19 @@ def _from_factorized(cls, values, original):
     # ------------------------------------------------------------------------
     # Must be a Sequence
     # ------------------------------------------------------------------------
+    @overload
+    def __getitem__(self, item: int | np.integer) -> Any:
+        ...
 
-    def __getitem__(self, item: PositionalIndexer) -> ExtensionArray | Any:
+    @overload
+    def __getitem__(
+        self: ExtensionArrayT, item: slice | np.ndarray | list[int]
+    ) -> ExtensionArrayT:
+        ...
+
+    def __getitem__(
+        self: ExtensionArrayT, item: PositionalIndexer
+    ) -> ExtensionArrayT | Any:
         """
         Select a subset of self.
 
@@ -306,6 +318,8 @@ def __getitem__(self, item: PositionalIndexer) -> ExtensionArray | Any:
 
             * ndarray: A 1-d boolean NumPy ndarray the same length as 'self'
 
+            * list[int]:  A list of int
+
         Returns
         -------
         item : scalar or ExtensionArray
@@ -736,7 +750,7 @@ def fillna(
             new_values = self.copy()
         return new_values
 
-    def dropna(self):
+    def dropna(self: ExtensionArrayT) -> ExtensionArrayT:
         """
         Return ExtensionArray without NA values.
 

pandas/core/arrays/categorical.py

@@ -11,6 +11,7 @@
     TypeVar,
     Union,
     cast,
+    overload,
 )
 from warnings import (
     catch_warnings,
@@ -34,6 +35,8 @@
     Dtype,
     NpDtype,
     Ordered,
+    PositionalIndexer2D,
+    PositionalIndexerTuple,
     Scalar,
     Shape,
     type_t,
@@ -2015,7 +2018,20 @@ def __repr__(self) -> str:
 
     # ------------------------------------------------------------------
 
-    def __getitem__(self, key):
+    @overload
+    def __getitem__(self, key: int | np.integer) -> object:
+        ...
+
+    @overload
+    def __getitem__(
+        self: CategoricalT,
+        key: slice | np.ndarray | list[int] | PositionalIndexerTuple,
+    ) -> CategoricalT:
+        ...
+
+    def __getitem__(
+        self: CategoricalT, key: PositionalIndexer2D
+    ) -> CategoricalT | object:
         """
         Return an item.
         """

pandas/core/arrays/datetimelike.py

@@ -48,6 +48,7 @@
     DtypeObj,
     NpDtype,
     PositionalIndexer2D,
+    PositionalIndexerTuple,
 )
 from pandas.compat.numpy import function as nv
 from pandas.errors import (
@@ -308,14 +309,26 @@ def __array__(self, dtype: NpDtype | None = None) -> np.ndarray:
             return np.array(list(self), dtype=object)
         return self._ndarray
 
+    @overload
+    def __getitem__(self, item: int | np.integer) -> DTScalarOrNaT:
+        ...
+
+    @overload
+    def __getitem__(
+        self: DatetimeLikeArrayT,
+        item: slice | np.ndarray | list[int] | PositionalIndexerTuple,
+    ) -> DatetimeLikeArrayT:
+        ...
+
     def __getitem__(
-        self, key: PositionalIndexer2D
-    ) -> DatetimeLikeArrayMixin | DTScalarOrNaT:
+        self: DatetimeLikeArrayT, key: PositionalIndexer2D
+    ) -> DatetimeLikeArrayT | DTScalarOrNaT:
         """
         This getitem defers to the underlying array, which by-definition can
         only handle list-likes, slices, and integer scalars
         """
-        result = super().__getitem__(key)
+        # Use cast as we know we will get back a DatetimeLikeArray
+        result = cast(DatetimeLikeArrayT, super().__getitem__(key))
         if lib.is_scalar(result):
             return result
 
@@ -1787,11 +1800,7 @@ def factorize(self, na_sentinel=-1, sort: bool = False):
             uniques = self.copy()  # TODO: copy or view?
             if sort and self.freq.n < 0:
                 codes = codes[::-1]
-                # TODO: overload __getitem__, a slice indexer returns same type as self
-                # error: Incompatible types in assignment (expression has type
-                # "Union[DatetimeLikeArrayMixin, Union[Any, Any]]", variable
-                # has type "TimelikeOps")
-                uniques = uniques[::-1]  # type: ignore[assignment]
+                uniques = uniques[::-1]
             return codes, uniques
         # FIXME: shouldn't get here; we are ignoring sort
         return super().factorize(na_sentinel=na_sentinel)

pandas/core/arrays/datetimes.py

@@ -8,7 +8,6 @@
 )
 from typing import (
     TYPE_CHECKING,
-    cast,
     overload,
 )
 import warnings
@@ -476,11 +475,9 @@ def _generate_range(
             index = cls._simple_new(arr, freq=None, dtype=dtype)
 
         if not left_closed and len(index) and index[0] == start:
-            # TODO: overload DatetimeLikeArrayMixin.__getitem__
-            index = cast(DatetimeArray, index[1:])
+            index = index[1:]
         if not right_closed and len(index) and index[-1] == end:
-            # TODO: overload DatetimeLikeArrayMixin.__getitem__
-            index = cast(DatetimeArray, index[:-1])
+            index = index[:-1]
 
         dtype = tz_to_dtype(tz)
         return cls._simple_new(index._ndarray, freq=freq, dtype=dtype)

pandas/core/arrays/interval.py

@@ -10,6 +10,7 @@
     Sequence,
     TypeVar,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -28,6 +29,7 @@
     ArrayLike,
     Dtype,
     NpDtype,
+    PositionalIndexer,
 )
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import Appender
@@ -629,7 +631,19 @@ def __iter__(self):
     def __len__(self) -> int:
         return len(self._left)
 
-    def __getitem__(self, key):
+    @overload
+    def __getitem__(self, key: int | np.integer) -> Interval:
+        ...
+
+    @overload
+    def __getitem__(
+        self: IntervalArrayT, key: slice | np.ndarray | list[int]
+    ) -> IntervalArrayT:
+        ...
+
+    def __getitem__(
+        self: IntervalArrayT, key: PositionalIndexer
+    ) -> IntervalArrayT | Interval:
         key = check_array_indexer(self, key)
         left = self._left[key]
         right = self._right[key]

pandas/core/arrays/masked.py

@@ -5,6 +5,7 @@
     Any,
     Sequence,
     TypeVar,
+    overload,
 )
 
 import numpy as np
@@ -135,7 +136,19 @@ def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False):
     def dtype(self) -> BaseMaskedDtype:
         raise AbstractMethodError(self)
 
-    def __getitem__(self, item: PositionalIndexer) -> BaseMaskedArray | Any:
+    @overload
+    def __getitem__(self, item: int | np.integer) -> Any:
+        ...
+
+    @overload
+    def __getitem__(
+        self: BaseMaskedArrayT, item: slice | np.ndarray | list[int]
+    ) -> BaseMaskedArrayT:
+        ...
+
+    def __getitem__(
+        self: BaseMaskedArrayT, item: PositionalIndexer
+    ) -> BaseMaskedArrayT | Any:
         if is_integer(item):
             if self._mask[item]:
                 return self.dtype.na_value

pandas/core/arrays/sparse/array.py

@@ -7,10 +7,13 @@
 import numbers
 import operator
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Sequence,
     TypeVar,
+    cast,
+    overload,
 )
 import warnings
 
@@ -27,6 +30,7 @@
 from pandas._typing import (
     Dtype,
     NpDtype,
+    PositionalIndexer,
     Scalar,
 )
 from pandas.compat.numpy import function as nv
@@ -77,6 +81,17 @@
 
 import pandas.io.formats.printing as printing
 
+# See https://github.com/python/typing/issues/684
+if TYPE_CHECKING:
+    from enum import Enum
+
+    class ellipsis(Enum):
+        Ellipsis = "..."
+
+    Ellipsis = ellipsis.Ellipsis
+else:
+    ellipsis = type(Ellipsis)
+
 # ----------------------------------------------------------------------------
 # Array
 
@@ -810,8 +825,21 @@ def value_counts(self, dropna: bool = True):
     # --------
     # Indexing
     # --------
+    @overload
+    def __getitem__(self, key: int | np.integer) -> Any:
+        ...
+
+    @overload
+    def __getitem__(
+        self: SparseArrayT,
+        key: slice | np.ndarray | list[int] | tuple[int | ellipsis, ...],
+    ) -> SparseArrayT:
+        ...
 
-    def __getitem__(self, key):
+    def __getitem__(
+        self: SparseArrayT,
+        key: PositionalIndexer | tuple[int | ellipsis, ...],
+    ) -> SparseArrayT | Any:
 
         if isinstance(key, tuple):
             if len(key) > 1:
@@ -821,6 +849,8 @@ def __getitem__(self, key):
                     key = key[:-1]
             if len(key) > 1:
                 raise IndexError("too many indices for array.")
+            if key[0] is Ellipsis:
+                raise ValueError("Cannot slice with Ellipsis")
             key = key[0]
 
         if is_integer(key):
@@ -849,7 +879,8 @@ def __getitem__(self, key):
             key = check_array_indexer(self, key)
 
             if com.is_bool_indexer(key):
-
+                # mypy doesn't know we have an array here
+                key = cast(np.ndarray, key)
                 return self.take(np.arange(len(key), dtype=np.int32)[key])
             elif hasattr(key, "__len__"):
                 return self.take(key)

pandas/core/arrays/string_arrow.py

@@ -6,6 +6,7 @@
     Any,
     Sequence,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -341,7 +342,19 @@ def _concat_same_type(cls, to_concat) -> ArrowStringArray:
             )
         )
 
-    def __getitem__(self, item: PositionalIndexer) -> Any:
+    @overload
+    def __getitem__(self, item: int | np.integer) -> str:
+        ...
+
+    @overload
+    def __getitem__(
+        self: ArrowStringArray, item: slice | np.ndarray | list[int]
+    ) -> ArrowStringArray:
+        ...
+
+    def __getitem__(
+        self: ArrowStringArray, item: PositionalIndexer
+    ) -> ArrowStringArray | str:
         """Select a subset of self.
 
         Parameters