TYP: Update type hints for ExtensionArray and ExtensionDtype

pandas/_testing/asserters.py

@@ -382,9 +382,13 @@ def _get_ilevel_values(index, level):
     # skip exact index checking when `check_categorical` is False
     if check_exact and check_categorical:
         if not left.equals(right):
-            diff = (
-                np.sum((left._values != right._values).astype(int)) * 100.0 / len(left)
-            )
+            #  error: Value of type variable "_Number" of "sum" cannot be
+            # "Union[ExtensionArray, ndarray, Any]"
+            thesum = np.sum(
+                (left._values != right._values).astype(int)
+            )  # type: ignore[type-var]
+            # error: Unsupported operand types for * ("ExtensionArray" and "float")
+            diff = thesum * 100.0 / len(left)  # type: ignore[operator]
             msg = f"{obj} values are different ({np.round(diff, 5)} %)"
             raise_assert_detail(obj, msg, left, right)
     else:

pandas/_typing.py

@@ -76,11 +76,17 @@
 
 ArrayLike = Union["ExtensionArray", np.ndarray]
 AnyArrayLike = Union[ArrayLike, "Index", "Series"]
-
+AnySequenceLike = Union[
+    "ExtensionArray",
+    "Index",
+    "Series",
+    Sequence[Any],
+    np.ndarray,
+]
 # scalars
 
 PythonScalar = Union[str, int, float, bool]
-DatetimeLikeScalar = TypeVar("DatetimeLikeScalar", "Period", "Timestamp", "Timedelta")
+DatetimeLikeScalar = Union["Period", "Timestamp", "Timedelta"]
 PandasScalar = Union["Period", "Timestamp", "Timedelta", "Interval"]
 Scalar = Union[PythonScalar, PandasScalar]
 

pandas/core/algorithms.py

@@ -30,7 +30,9 @@
 )
 from pandas._typing import (
     AnyArrayLike,
+    AnySequenceLike,
     ArrayLike,
+    Dtype,
     DtypeObj,
     FrameOrSeriesUnion,
 )
@@ -216,7 +218,7 @@ def _ensure_data(values: ArrayLike) -> Tuple[np.ndarray, DtypeObj]:
 
 
 def _reconstruct_data(
-    values: ArrayLike, dtype: DtypeObj, original: AnyArrayLike
+    values: ArrayLike, dtype: Dtype, original: AnyArrayLike
 ) -> ArrayLike:
     """
     reverse of _ensure_data
@@ -446,7 +448,7 @@ def unique(values):
 unique1d = unique
 
 
-def isin(comps: AnyArrayLike, values: AnyArrayLike) -> np.ndarray:
+def isin(comps: AnySequenceLike, values: AnySequenceLike) -> np.ndarray:
     """
     Compute the isin boolean array.
 
@@ -482,9 +484,11 @@ def isin(comps: AnyArrayLike, values: AnyArrayLike) -> np.ndarray:
     comps = _ensure_arraylike(comps)
     comps = extract_array(comps, extract_numpy=True)
     if is_extension_array_dtype(comps.dtype):
-        # error: Incompatible return value type (got "Series", expected "ndarray")
-        # error: Item "ndarray" of "Union[Any, ndarray]" has no attribute "isin"
-        return comps.isin(values)  # type: ignore[return-value,union-attr]
+        #  error: Argument 1 to "isin" of "ExtensionArray" has incompatible type
+        # "Union[Any, ExtensionArray, ndarray]"; expected "Sequence[Any]"
+        # error: Item "ndarray" of "Union[Any, ExtensionArray, ndarray]" has no
+        # attribute "isin"
+        return comps.isin(values)  # type: ignore[arg-type, union-attr]
 
     elif needs_i8_conversion(comps.dtype):
         # Dispatch to DatetimeLikeArrayMixin.isin

pandas/core/array_algos/putmask.py

@@ -191,7 +191,7 @@ def extract_bool_array(mask: ArrayLike) -> np.ndarray:
         # We could have BooleanArray, Sparse[bool], ...
         #  Except for BooleanArray, this is equivalent to just
         #  np.asarray(mask, dtype=bool)
-        mask = mask.to_numpy(dtype=bool, na_value=False)
+        mask = mask.to_numpy(dtype=np.dtype(bool), na_value=False)
 
     mask = np.asarray(mask, dtype=bool)
     return mask

pandas/core/arrays/_mixins.py

@@ -3,11 +3,13 @@
 from functools import wraps
 from typing import (
     Any,
+    List,
     Optional,
     Sequence,
     Type,
     TypeVar,
     Union,
+    overload,
 )
 
 import numpy as np
@@ -253,8 +255,19 @@ def __setitem__(self, key, value):
     def _validate_setitem_value(self, value):
         return value
 
+    @overload
+    def __getitem__(self: NDArrayBackedExtensionArrayT, key: int) -> Any:
+        ...
+
+    @overload
     def __getitem__(
-        self: NDArrayBackedExtensionArrayT, key: Union[int, slice, np.ndarray]
+        self: NDArrayBackedExtensionArrayT, key: Union[slice, np.ndarray, List[Any]]
+    ) -> NDArrayBackedExtensionArrayT:
+        ...
+
+    def __getitem__(
+        self: NDArrayBackedExtensionArrayT,
+        key: Union[int, slice, np.ndarray, List[Any]],
     ) -> Union[NDArrayBackedExtensionArrayT, Any]:
         if lib.is_integer(key):
             # fast-path