TYP: Typing for ExtensionArray.__getitem__

pandas/core/arrays/_mixins.py

@@ -6,6 +6,7 @@
     Sequence,
     TypeVar,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -276,10 +277,20 @@ def __setitem__(self, key, value):
     def _validate_setitem_value(self, value):
         return value
 
+    @overload
+    def __getitem__(self, item: int | np.integer) -> Any:
+        ...
+
+    @overload
     def __getitem__(
-        self: NDArrayBackedExtensionArrayT,
+        self: NDArrayBackedExtensionArray, item: slice | np.ndarray | Sequence[int]
+    ) -> NDArrayBackedExtensionArray:
+        ...
+
+    def __getitem__(
+        self: NDArrayBackedExtensionArray,
         key: PositionalIndexer2D,
-    ) -> NDArrayBackedExtensionArrayT | Any:
+    ) -> NDArrayBackedExtensionArray | Any:
         if lib.is_integer(key):
             # fast-path
             result = self._ndarray[key]

pandas/core/arrays/base.py

@@ -16,6 +16,7 @@
     Sequence,
     TypeVar,
     cast,
+    overload,
 )
 
 import numpy as np
@@ -24,7 +25,7 @@
 from pandas._typing import (
     ArrayLike,
     Dtype,
-    PositionalIndexer,
+    PositionalIndexer2D,
     Shape,
 )
 from pandas.compat import set_function_name
@@ -288,8 +289,15 @@ 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, item: slice | np.ndarray | Sequence[int]) -> ExtensionArray:
+        ...
+
+    def __getitem__(self, item: PositionalIndexer2D) -> ExtensionArray | Any:
         """
         Select a subset of self.
 
@@ -729,7 +737,7 @@ def fillna(self, value=None, method=None, limit=None):
             new_values = self.copy()
         return new_values
 
-    def dropna(self):
+    def dropna(self: ExtensionArrayT) -> ExtensionArrayT:
         """
         Return ExtensionArray without NA values.
 

pandas/core/arrays/datetimelike.py

@@ -309,14 +309,28 @@ 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, item: slice | np.ndarray | Sequence[int]
+    ) -> DatetimeLikeArrayMixin:
+        ...
+
     def __getitem__(
         self, key: PositionalIndexer2D
     ) -> DatetimeLikeArrayMixin | DTScalarOrNaT:
         """
         This getitem defers to the underlying array, which by-definition can
         only handle list-likes, slices, and integer scalars
         """
-        result = super().__getitem__(key)
+        #  error: Invalid index type "Union[Union[int, integer[Any], slice,
+        # Sequence[int], ndarray], Tuple[Union[int, integer[Any], slice, Sequence[int],
+        # ndarray], Union[int, integer[Any], slice, Sequence[int], ndarray]]]" for
+        # "DatetimeLikeArrayMixin"; expected type "Union[int, integer[Any]]"
+        result = super().__getitem__(key)  # type: ignore[index]
         if lib.is_scalar(result):
             return result
 
@@ -1778,11 +1792,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/indexes/base.py

@@ -4595,11 +4595,7 @@ def __getitem__(self, key):
 
         result = getitem(key)
         if not is_scalar(result):
-            # error: Argument 1 to "ndim" has incompatible type "Union[ExtensionArray,
-            # Any]"; expected "Union[Union[int, float, complex, str, bytes, generic],
-            # Sequence[Union[int, float, complex, str, bytes, generic]],
-            # Sequence[Sequence[Any]], _SupportsArray]"
-            if np.ndim(result) > 1:  # type: ignore[arg-type]
+            if np.ndim(result) > 1:
                 deprecate_ndim_indexing(result)
                 return result
             # NB: Using _constructor._simple_new would break if MultiIndex

pandas/core/missing.py

@@ -421,11 +421,7 @@ def interpolate_1d(
     if method in NP_METHODS:
         # np.interp requires sorted X values, #21037
 
-        # error: Argument 1 to "argsort" has incompatible type "Union[ExtensionArray,
-        # Any]"; expected "Union[Union[int, float, complex, str, bytes, generic],
-        # Sequence[Union[int, float, complex, str, bytes, generic]],
-        # Sequence[Sequence[Any]], _SupportsArray]"
-        indexer = np.argsort(inds[valid])  # type: ignore[arg-type]
+        indexer = np.argsort(inds[valid])
         result[invalid] = np.interp(
             inds[invalid], inds[valid][indexer], yvalues[valid][indexer]
         )