TYP: __getitem__ method of EA (#37898)

Author: simonjayhawkins

pandas/core/arrays/_mixins.py

@@ -1,4 +1,6 @@
-from typing import Any, Optional, Sequence, Type, TypeVar
+from __future__ import annotations
+
+from typing import Any, Optional, Sequence, Type, TypeVar, Union
 
 import numpy as np
 
@@ -212,7 +214,9 @@ def __setitem__(self, key, value):
     def _validate_setitem_value(self, value):
         return value
 
-    def __getitem__(self, key):
+    def __getitem__(
+        self: NDArrayBackedExtensionArrayT, key: Union[int, slice, np.ndarray]
+    ) -> Union[NDArrayBackedExtensionArrayT, Any]:
         if lib.is_integer(key):
             # fast-path
             result = self._ndarray[key]

pandas/core/arrays/base.py

@@ -6,6 +6,8 @@
    This is an experimental API and subject to breaking changes
    without warning.
 """
+from __future__ import annotations
+
 import operator
 from typing import (
     Any,
@@ -254,8 +256,9 @@ def _from_factorized(cls, values, original):
     # Must be a Sequence
     # ------------------------------------------------------------------------
 
-    def __getitem__(self, item):
-        # type (Any) -> Any
+    def __getitem__(
+        self, item: Union[int, slice, np.ndarray]
+    ) -> Union[ExtensionArray, Any]:
         """
         Select a subset of self.
 
@@ -661,7 +664,7 @@ def dropna(self):
         """
         return self[~self.isna()]
 
-    def shift(self, periods: int = 1, fill_value: object = None) -> "ExtensionArray":
+    def shift(self, periods: int = 1, fill_value: object = None) -> ExtensionArray:
         """
         Shift values by desired number.
 
@@ -831,7 +834,7 @@ def _values_for_factorize(self) -> Tuple[np.ndarray, Any]:
         """
         return self.astype(object), np.nan
 
-    def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, "ExtensionArray"]:
+    def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, ExtensionArray]:
         """
         Encode the extension array as an enumerated type.
 
@@ -940,7 +943,7 @@ def take(
         *,
         allow_fill: bool = False,
         fill_value: Any = None,
-    ) -> "ExtensionArray":
+    ) -> ExtensionArray:
         """
         Take elements from an array.
 
@@ -1109,7 +1112,7 @@ def _formatter(self, boxed: bool = False) -> Callable[[Any], Optional[str]]:
     # Reshaping
     # ------------------------------------------------------------------------
 
-    def transpose(self, *axes) -> "ExtensionArray":
+    def transpose(self, *axes) -> ExtensionArray:
         """
         Return a transposed view on this array.
 
@@ -1119,10 +1122,10 @@ def transpose(self, *axes) -> "ExtensionArray":
         return self[:]
 
     @property
-    def T(self) -> "ExtensionArray":
+    def T(self) -> ExtensionArray:
         return self.transpose()
 
-    def ravel(self, order="C") -> "ExtensionArray":
+    def ravel(self, order="C") -> ExtensionArray:
         """
         Return a flattened view on this array.
 

pandas/core/arrays/datetimelike.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from datetime import datetime, timedelta
 import operator
 from typing import (
@@ -264,7 +266,9 @@ def __array__(self, dtype=None) -> np.ndarray:
             return np.array(list(self), dtype=object)
         return self._ndarray
 
-    def __getitem__(self, key):
+    def __getitem__(
+        self, key: Union[int, slice, np.ndarray]
+    ) -> Union[DatetimeLikeArrayMixin, DTScalarOrNaT]:
         """
         This getitem defers to the underlying array, which by-definition can
         only handle list-likes, slices, and integer scalars

pandas/core/arrays/datetimes.py

@@ -1,5 +1,5 @@
 from datetime import datetime, time, timedelta, tzinfo
-from typing import Optional, Union
+from typing import Optional, Union, cast
 import warnings
 
 import numpy as np
@@ -444,9 +444,11 @@ def _generate_range(
             )
 
         if not left_closed and len(index) and index[0] == start:
-            index = index[1:]
+            # TODO: overload DatetimeLikeArrayMixin.__getitem__
+            index = cast(DatetimeArray, index[1:])
         if not right_closed and len(index) and index[-1] == end:
-            index = index[:-1]
+            # TODO: overload DatetimeLikeArrayMixin.__getitem__
+            index = cast(DatetimeArray, index[:-1])
 
         dtype = tz_to_dtype(tz)
         return cls._simple_new(index.asi8, freq=freq, dtype=dtype)

pandas/core/arrays/masked.py

@@ -1,4 +1,6 @@
-from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Type, TypeVar
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Tuple, Type, TypeVar, Union
 
 import numpy as np
 
@@ -56,7 +58,7 @@ def itemsize(self) -> int:
         return self.numpy_dtype.itemsize
 
     @classmethod
-    def construct_array_type(cls) -> Type["BaseMaskedArray"]:
+    def construct_array_type(cls) -> Type[BaseMaskedArray]:
         """
         Return the array type associated with this dtype.
 
@@ -100,7 +102,9 @@ def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False):
     def dtype(self) -> BaseMaskedDtype:
         raise AbstractMethodError(self)
 
-    def __getitem__(self, item):
+    def __getitem__(
+        self, item: Union[int, slice, np.ndarray]
+    ) -> Union[BaseMaskedArray, Any]:
         if is_integer(item):
             if self._mask[item]:
                 return self.dtype.na_value

pandas/core/groupby/groupby.py

@@ -1671,10 +1671,10 @@ def first(self, numeric_only: bool = False, min_count: int = -1):
         def first_compat(obj: FrameOrSeries, axis: int = 0):
             def first(x: Series):
                 """Helper function for first item that isn't NA."""
-                x = x.array[notna(x.array)]
-                if len(x) == 0:
+                arr = x.array[notna(x.array)]
+                if not len(arr):
                     return np.nan
-                return x[0]
+                return arr[0]
 
             if isinstance(obj, DataFrame):
                 return obj.apply(first, axis=axis)
@@ -1695,10 +1695,10 @@ def last(self, numeric_only: bool = False, min_count: int = -1):
         def last_compat(obj: FrameOrSeries, axis: int = 0):
             def last(x: Series):
                 """Helper function for last item that isn't NA."""
-                x = x.array[notna(x.array)]
-                if len(x) == 0:
+                arr = x.array[notna(x.array)]
+                if not len(arr):
                     return np.nan
-                return x[-1]
+                return arr[-1]
 
             if isinstance(obj, DataFrame):
                 return obj.apply(last, axis=axis)

pandas/core/indexes/base.py

@@ -3222,8 +3222,14 @@ def _get_nearest_indexer(self, target: "Index", limit, tolerance) -> np.ndarray:
         right_indexer = self.get_indexer(target, "backfill", limit=limit)
 
         target_values = target._values
-        left_distances = np.abs(self._values[left_indexer] - target_values)
-        right_distances = np.abs(self._values[right_indexer] - target_values)
+        # error: Unsupported left operand type for - ("ExtensionArray")
+        left_distances = np.abs(
+            self._values[left_indexer] - target_values  # type: ignore[operator]
+        )
+        # error: Unsupported left operand type for - ("ExtensionArray")
+        right_distances = np.abs(
+            self._values[right_indexer] - target_values  # type: ignore[operator]
+        )
 
         op = operator.lt if self.is_monotonic_increasing else operator.le
         indexer = np.where(
@@ -3242,7 +3248,8 @@ def _filter_indexer_tolerance(
         indexer: np.ndarray,
         tolerance,
     ) -> np.ndarray:
-        distance = abs(self._values[indexer] - target)
+        # error: Unsupported left operand type for - ("ExtensionArray")
+        distance = abs(self._values[indexer] - target)  # type: ignore[operator]
         indexer = np.where(distance <= tolerance, indexer, -1)
         return indexer
 
@@ -3446,6 +3453,7 @@ def reindex(self, target, method=None, level=None, limit=None, tolerance=None):
         target = ensure_has_len(target)  # target may be an iterator
 
         if not isinstance(target, Index) and len(target) == 0:
+            values: Union[range, ExtensionArray, np.ndarray]
             if isinstance(self, ABCRangeIndex):
                 values = range(0)
             else:
@@ -4538,8 +4546,9 @@ def asof_locs(self, where: "Index", mask) -> np.ndarray:
 
         result = np.arange(len(self))[mask].take(locs)
 
-        first = mask.argmax()
-        result[(locs == 0) & (where._values < self._values[first])] = -1
+        # TODO: overload return type of ExtensionArray.__getitem__
+        first_value = cast(Any, self._values[mask.argmax()])
+        result[(locs == 0) & (where._values < first_value)] = -1
 
         return result
 

pandas/core/indexes/period.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timedelta
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 
@@ -673,7 +673,10 @@ def difference(self, other, sort=None):
 
         if self.equals(other):
             # pass an empty PeriodArray with the appropriate dtype
-            return type(self)._simple_new(self._data[:0], name=self.name)
+
+            # TODO: overload DatetimeLikeArrayMixin.__getitem__
+            values = cast(PeriodArray, self._data[:0])
+            return type(self)._simple_new(values, name=self.name)
 
         if is_object_dtype(other):
             return self.astype(object).difference(other).astype(self.dtype)