TYP: Typing for ExtensionArray.__getitem__

pandas/_typing.py

@@ -204,7 +204,7 @@
 # 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]
-]
+# Using List[int] here rather than Sequence[int] to disallow tuples.
+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,
 )
@@ -185,6 +187,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: NDArrayBackedExtensionArrayT,
+        key: slice | np.ndarray | list[int] | PositionalIndexerTuple,
+    ) -> NDArrayBackedExtensionArrayT:
+        ...
+
     def __getitem__(
         self: NDArrayBackedExtensionArrayT,
         key: PositionalIndexer2D,

pandas/core/arrays/base.py

@@ -18,6 +18,7 @@
     Sequence,
     TypeVar,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -290,8 +291,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.
 
@@ -305,6 +317,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
@@ -735,7 +749,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

@@ -6,11 +6,13 @@
 from shutil import get_terminal_size
 from typing import (
     TYPE_CHECKING,
+    Any,
     Hashable,
     Sequence,
     TypeVar,
     Union,
     cast,
+    overload,
 )
 from warnings import (
     catch_warnings,
@@ -35,6 +37,8 @@
     Dtype,
     NpDtype,
     Ordered,
+    PositionalIndexer2D,
+    PositionalIndexerTuple,
     Scalar,
     Shape,
     npt,
@@ -2003,7 +2007,18 @@ def __repr__(self) -> str:
 
     # ------------------------------------------------------------------
 
-    def __getitem__(self, key):
+    @overload
+    def __getitem__(self, key: int | np.integer) -> Any:
+        ...
+
+    @overload
+    def __getitem__(
+        self: CategoricalT,
+        key: slice | np.ndarray | list[int] | PositionalIndexerTuple,
+    ) -> CategoricalT:
+        ...
+
+    def __getitem__(self: CategoricalT, key: PositionalIndexer2D) -> CategoricalT | Any:
         """
         Return an item.
         """

pandas/core/arrays/datetimelike.py

@@ -49,6 +49,7 @@
     DtypeObj,
     NpDtype,
     PositionalIndexer2D,
+    PositionalIndexerTuple,
 )
 from pandas.compat.numpy import function as nv
 from pandas.errors import (
@@ -312,17 +313,33 @@ 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, key: PositionalIndexer2D
-    ) -> DatetimeLikeArrayMixin | DTScalarOrNaT:
+        self: DatetimeLikeArrayT,
+        item: slice | np.ndarray | list[int] | PositionalIndexerTuple,
+    ) -> DatetimeLikeArrayT:
+        ...
+
+    def __getitem__(
+        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 or DTScalar
+        result = cast(
+            Union[DatetimeLikeArrayT, DTScalarOrNaT], super().__getitem__(key)
+        )
         if lib.is_scalar(result):
             return result
-
+        else:
+            # At this point we know the result is an array.
+            result = cast(DatetimeLikeArrayT, result)
         result._freq = self._get_getitem_freq(key)
         return result
 
@@ -1767,11 +1784,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

@@ -9,7 +9,6 @@
 from typing import (
     TYPE_CHECKING,
     Literal,
-    cast,
     overload,
 )
 import warnings
@@ -475,11 +474,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

@@ -9,7 +9,9 @@
 from typing import (
     Sequence,
     TypeVar,
+    Union,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -28,6 +30,7 @@
     ArrayLike,
     Dtype,
     NpDtype,
+    PositionalIndexer,
 )
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import Appender
@@ -86,6 +89,7 @@
 )
 
 IntervalArrayT = TypeVar("IntervalArrayT", bound="IntervalArray")
+IntervalOrNA = Union[Interval, float]
 
 _interval_shared_docs: dict[str, str] = {}
 
@@ -629,7 +633,19 @@ def __iter__(self):
     def __len__(self) -> int:
         return len(self._left)
 
-    def __getitem__(self, key):
+    @overload
+    def __getitem__(self, key: int | np.integer) -> IntervalOrNA:
+        ...
+
+    @overload
+    def __getitem__(
+        self: IntervalArrayT, key: slice | np.ndarray | list[int]
+    ) -> IntervalArrayT:
+        ...
+
+    def __getitem__(
+        self: IntervalArrayT, key: PositionalIndexer
+    ) -> IntervalArrayT | IntervalOrNA:
         key = check_array_indexer(self, key)
         left = self._left[key]
         right = self._right[key]
@@ -1627,10 +1643,11 @@ def _from_combined(self, combined: np.ndarray) -> IntervalArray:
         return self._shallow_copy(left=new_left, right=new_right)
 
     def unique(self) -> IntervalArray:
-        # Invalid index type "Tuple[slice, int]" for "Union[ExtensionArray,
-        # ndarray[Any, Any]]"; expected type "Union[int, integer[Any], slice,
-        # Sequence[int], ndarray[Any, Any]]"
-        nc = unique(self._combined.view("complex128")[:, 0])  # type: ignore[index]
+        # No overload variant of "__getitem__" of "ExtensionArray" matches argument
+        # type "Tuple[slice, int]"
+        nc = unique(
+            self._combined.view("complex128")[:, 0]  # type: ignore[call-overload]
+        )
         nc = nc[:, None]
         return self._from_combined(nc)
 

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
 
@@ -806,8 +821,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:
@@ -817,6 +845,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):
@@ -845,7 +875,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)