TYP: ExtensionArray.take accept np.ndarray

pandas/_typing.py

@@ -219,6 +219,10 @@
 PositionalIndexer = Union[ScalarIndexer, SequenceIndexer]
 PositionalIndexerTuple = Tuple[PositionalIndexer, PositionalIndexer]
 PositionalIndexer2D = Union[PositionalIndexer, PositionalIndexerTuple]
+if TYPE_CHECKING:
+    TakeIndexer = Union[Sequence[int], Sequence[np.integer], npt.NDArray[np.integer]]
+else:
+    TakeIndexer = Any
 
 # Windowing rank methods
 WindowingRankType = Literal["average", "min", "max"]

pandas/core/algorithms.py

@@ -28,6 +28,7 @@
     ArrayLike,
     DtypeObj,
     Scalar,
+    TakeIndexer,
     npt,
 )
 from pandas.util._decorators import doc
@@ -1431,7 +1432,11 @@ def get_indexer(current_indexer, other_indexer):
 
 
 def take(
-    arr, indices: np.ndarray, axis: int = 0, allow_fill: bool = False, fill_value=None
+    arr,
+    indices: TakeIndexer,
+    axis: int = 0,
+    allow_fill: bool = False,
+    fill_value=None,
 ):
     """
     Take elements from an array.
@@ -1441,7 +1446,7 @@ def take(
     arr : array-like or scalar value
         Non array-likes (sequences/scalars without a dtype) are coerced
         to an ndarray.
-    indices : sequence of integers
+    indices : sequence of int or one-dimensional np.ndarray of int
         Indices to be taken.
     axis : int, default 0
         The axis over which to select values.

pandas/core/array_algos/take.py

@@ -178,14 +178,7 @@ def take_1d(
     """
     if not isinstance(arr, np.ndarray):
         # ExtensionArray -> dispatch to their method
-
-        # error: Argument 1 to "take" of "ExtensionArray" has incompatible type
-        # "ndarray"; expected "Sequence[int]"
-        return arr.take(
-            indexer,  # type: ignore[arg-type]
-            fill_value=fill_value,
-            allow_fill=allow_fill,
-        )
+        return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill)
 
     if not allow_fill:
         return arr.take(indexer)

pandas/core/arrays/_mixins.py

@@ -22,6 +22,7 @@
     ScalarIndexer,
     SequenceIndexer,
     Shape,
+    TakeIndexer,
     npt,
     type_t,
 )
@@ -101,7 +102,7 @@ def _validate_scalar(self, value):
 
     def take(
         self: NDArrayBackedExtensionArrayT,
-        indices: Sequence[int],
+        indices: TakeIndexer,
         *,
         allow_fill: bool = False,
         fill_value: Any = None,
@@ -112,9 +113,7 @@ def take(
 
         new_data = take(
             self._ndarray,
-            # error: Argument 2 to "take" has incompatible type "Sequence[int]";
-            # expected "ndarray"
-            indices,  # type: ignore[arg-type]
+            indices,
             allow_fill=allow_fill,
             fill_value=fill_value,
             axis=axis,

pandas/core/arrays/base.py

@@ -33,6 +33,7 @@
     ScalarIndexer,
     SequenceIndexer,
     Shape,
+    TakeIndexer,
     npt,
 )
 from pandas.compat import set_function_name
@@ -1076,7 +1077,7 @@ def repeat(self, repeats: int | Sequence[int], axis: int | None = None):
 
     def take(
         self: ExtensionArrayT,
-        indices: Sequence[int],
+        indices: TakeIndexer,
         *,
         allow_fill: bool = False,
         fill_value: Any = None,
@@ -1086,7 +1087,7 @@ def take(
 
         Parameters
         ----------
-        indices : sequence of int
+        indices : sequence of int or one-dimensional np.ndarray of int
             Indices to be taken.
         allow_fill : bool, default False
             How to handle negative values in `indices`.

pandas/core/arrays/string_arrow.py

@@ -5,7 +5,6 @@
 from typing import (
     TYPE_CHECKING,
     Any,
-    Sequence,
     Union,
     cast,
     overload,
@@ -24,6 +23,7 @@
     Scalar,
     ScalarIndexer,
     SequenceIndexer,
+    TakeIndexer,
     npt,
 )
 from pandas.compat import (
@@ -307,9 +307,7 @@ def __getitem__(
             if not len(item):
                 return type(self)(pa.chunked_array([], type=pa.string()))
             elif is_integer_dtype(item.dtype):
-                # error: Argument 1 to "take" of "ArrowStringArray" has incompatible
-                # type "ndarray"; expected "Sequence[int]"
-                return self.take(item)  # type: ignore[arg-type]
+                return self.take(item)
             elif is_bool_dtype(item.dtype):
                 return type(self)(self._data.filter(item))
             else:
@@ -513,14 +511,17 @@ def __setitem__(self, key: int | slice | np.ndarray, value: Any) -> None:
                 self[k] = v
 
     def take(
-        self, indices: Sequence[int], allow_fill: bool = False, fill_value: Any = None
+        self,
+        indices: TakeIndexer,
+        allow_fill: bool = False,
+        fill_value: Any = None,
     ):
         """
         Take elements from an array.
 
         Parameters
         ----------
-        indices : sequence of int
+        indices : sequence of int or one-dimensional np.ndarray of int
             Indices to be taken.
         allow_fill : bool, default False
             How to handle negative values in `indices`.

pandas/core/indexes/base.py

@@ -4318,12 +4318,8 @@ def _join_non_unique(
         )
         mask = left_idx == -1
 
-        # error: Argument 1 to "take" of "ExtensionArray" has incompatible
-        # type "ndarray[Any, dtype[signedinteger[Any]]]"; expected "Sequence[int]"
-        join_array = self._values.take(left_idx)  # type: ignore[arg-type]
-        # error: Argument 1 to "take" of "ExtensionArray" has incompatible type
-        # "ndarray[Any, dtype[signedinteger[Any]]]"; expected "Sequence[int]"
-        right = other._values.take(right_idx)  # type: ignore[arg-type]
+        join_array = self._values.take(left_idx)
+        right = other._values.take(right_idx)
 
         if isinstance(join_array, np.ndarray):
             np.putmask(join_array, mask, right)