TYP: Typing for ExtensionArray.__getitem__ (#41258)

Author: Dr-Irv

pandas/_typing.py

@@ -206,10 +206,16 @@
 # indexing
 # PositionalIndexer -> valid 1D positional indexer, e.g. can pass
 # to ndarray.__getitem__
+# ScalarIndexer is for a single value as the index
+# SequenceIndexer is for list like or slices (but not tuples)
+# PositionalIndexerTuple is extends the PositionalIndexer for 2D arrays
+# These are used in various __getitem__ overloads
 # 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.
+ScalarIndexer = Union[int, np.integer]
+SequenceIndexer = Union[slice, List[int], np.ndarray]
+PositionalIndexer = Union[ScalarIndexer, SequenceIndexer]
+PositionalIndexerTuple = Tuple[PositionalIndexer, PositionalIndexer]
+PositionalIndexer2D = Union[PositionalIndexer, PositionalIndexerTuple]

pandas/core/arrays/_mixins.py

@@ -4,9 +4,11 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Literal,
     Sequence,
     TypeVar,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -16,6 +18,9 @@
 from pandas._typing import (
     F,
     PositionalIndexer2D,
+    PositionalIndexerTuple,
+    ScalarIndexer,
+    SequenceIndexer,
     Shape,
     npt,
     type_t,
@@ -48,7 +53,6 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Literal
 
     from pandas._typing import (
         NumpySorter,
@@ -205,6 +209,17 @@ def __setitem__(self, key, value):
     def _validate_setitem_value(self, value):
         return value
 
+    @overload
+    def __getitem__(self, key: ScalarIndexer) -> Any:
+        ...
+
+    @overload
+    def __getitem__(
+        self: NDArrayBackedExtensionArrayT,
+        key: SequenceIndexer | PositionalIndexerTuple,
+    ) -> NDArrayBackedExtensionArrayT:
+        ...
+
     def __getitem__(
         self: NDArrayBackedExtensionArrayT,
         key: PositionalIndexer2D,

pandas/core/arrays/base.py

@@ -30,6 +30,8 @@
     Dtype,
     FillnaOptions,
     PositionalIndexer,
+    ScalarIndexer,
+    SequenceIndexer,
     Shape,
     npt,
 )
@@ -298,8 +300,17 @@ def _from_factorized(cls, values, original):
     # ------------------------------------------------------------------------
     # Must be a Sequence
     # ------------------------------------------------------------------------
+    @overload
+    def __getitem__(self, item: ScalarIndexer) -> Any:
+        ...
+
+    @overload
+    def __getitem__(self: ExtensionArrayT, item: SequenceIndexer) -> ExtensionArrayT:
+        ...
 
-    def __getitem__(self, item: PositionalIndexer) -> ExtensionArray | Any:
+    def __getitem__(
+        self: ExtensionArrayT, item: PositionalIndexer
+    ) -> ExtensionArrayT | Any:
         """
         Select a subset of self.
 
@@ -313,6 +324,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
@@ -761,7 +774,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,6 +6,7 @@
 from shutil import get_terminal_size
 from typing import (
     TYPE_CHECKING,
+    Any,
     Hashable,
     Sequence,
     TypeVar,
@@ -37,7 +38,11 @@
     Dtype,
     NpDtype,
     Ordered,
+    PositionalIndexer2D,
+    PositionalIndexerTuple,
     Scalar,
+    ScalarIndexer,
+    SequenceIndexer,
     Shape,
     npt,
     type_t,
@@ -2017,7 +2022,18 @@ def __repr__(self) -> str:
 
     # ------------------------------------------------------------------
 
-    def __getitem__(self, key):
+    @overload
+    def __getitem__(self, key: ScalarIndexer) -> Any:
+        ...
+
+    @overload
+    def __getitem__(
+        self: CategoricalT,
+        key: SequenceIndexer | PositionalIndexerTuple,
+    ) -> CategoricalT:
+        ...
+
+    def __getitem__(self: CategoricalT, key: PositionalIndexer2D) -> CategoricalT | Any:
         """
         Return an item.
         """

pandas/core/arrays/datetimelike.py

@@ -49,6 +49,9 @@
     DtypeObj,
     NpDtype,
     PositionalIndexer2D,
+    PositionalIndexerTuple,
+    ScalarIndexer,
+    SequenceIndexer,
     npt,
 )
 from pandas.compat.numpy import function as nv
@@ -313,17 +316,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: ScalarIndexer) -> DTScalarOrNaT:
+        ...
+
+    @overload
     def __getitem__(
-        self, key: PositionalIndexer2D
-    ) -> DatetimeLikeArrayMixin | DTScalarOrNaT:
+        self: DatetimeLikeArrayT,
+        item: SequenceIndexer | 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
 
@@ -1768,11 +1787,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
@@ -478,11 +477,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
@@ -31,6 +33,9 @@
     ArrayLike,
     Dtype,
     NpDtype,
+    PositionalIndexer,
+    ScalarIndexer,
+    SequenceIndexer,
 )
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import Appender
@@ -89,6 +94,7 @@
 )
 
 IntervalArrayT = TypeVar("IntervalArrayT", bound="IntervalArray")
+IntervalOrNA = Union[Interval, float]
 
 _interval_shared_docs: dict[str, str] = {}
 
@@ -635,7 +641,17 @@ def __iter__(self):
     def __len__(self) -> int:
         return len(self._left)
 
-    def __getitem__(self, key):
+    @overload
+    def __getitem__(self, key: ScalarIndexer) -> IntervalOrNA:
+        ...
+
+    @overload
+    def __getitem__(self: IntervalArrayT, key: SequenceIndexer) -> IntervalArrayT:
+        ...
+
+    def __getitem__(
+        self: IntervalArrayT, key: PositionalIndexer
+    ) -> IntervalArrayT | IntervalOrNA:
         key = check_array_indexer(self, key)
         left = self._left[key]
         right = self._right[key]
@@ -1633,10 +1649,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

@@ -20,6 +20,8 @@
     NpDtype,
     PositionalIndexer,
     Scalar,
+    ScalarIndexer,
+    SequenceIndexer,
     npt,
     type_t,
 )
@@ -139,7 +141,17 @@ 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: ScalarIndexer) -> Any:
+        ...
+
+    @overload
+    def __getitem__(self: BaseMaskedArrayT, item: SequenceIndexer) -> 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/period.py

@@ -6,6 +6,7 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    Literal,
     Sequence,
 )
 
@@ -76,7 +77,6 @@
 import pandas.core.common as com
 
 if TYPE_CHECKING:
-    from typing import Literal
 
     from pandas._typing import (
         NumpySorter,