Repository: pandas-dev/pandas
Pull request: TYP: ExtensionArray.take accept np.ndarray
- [x] closes #43391 (np.ndarray is never a Sequence)
- [x] tests added / passed
- [x] Ensure all linting tests pass, see [here](https://pandas.pydata.org/pandas-docs/dev/development/contributing.html#code-standards) for how to run them
- [ ] whatsnew entry

Problem statement:
TYP:  ensure_platform_int probably needs a dtype annotation
`ensure_platform_int` is annotated as returning a `np.ndarray` but its output is used in places that require a `Sequence[int]`. I assume this can be fixed by adding a dtype to the `np.ndarray`.

pandas/core/groupby/ops.py:858: error: Argument 1 to "take" of "ExtensionArray" has incompatible type "ndarray[Any, Any]"; expected "Sequence[int]"  [arg-type]

Hints from issue discussion:

Commit message: Update pandas/_typing.py

Co-authored-by: Simon Hawkins <simonjayhawkins@gmail.com>

Patch to review:
diff --git a/pandas/_typing.py b/pandas/_typing.py
index 9ed31dc3738f3..7d2ac6571dd8d 100644
--- a/pandas/_typing.py
+++ b/pandas/_typing.py
@@ -216,6 +216,7 @@
 # Using List[int] here rather than Sequence[int] to disallow tuples.
 ScalarIndexer = Union[int, np.integer]
 SequenceIndexer = Union[slice, List[int], np.ndarray]
+TakeIndexer = Union[Sequence[ScalarIndexer], npt.NDArray[np.integer]]
 PositionalIndexer = Union[ScalarIndexer, SequenceIndexer]
 PositionalIndexerTuple = Tuple[PositionalIndexer, PositionalIndexer]
 PositionalIndexer2D = Union[PositionalIndexer, PositionalIndexerTuple]
diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py
index bb4e0dff0f4c7..b8fb646001523 100644
--- a/pandas/core/algorithms.py
+++ b/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.
diff --git a/pandas/core/array_algos/take.py b/pandas/core/array_algos/take.py
index 201e177d8bb10..6d350cfa2c1d6 100644
--- a/pandas/core/array_algos/take.py
+++ b/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)
diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py
index 4c7ccc2f16477..22062f7d165f9 100644
--- a/pandas/core/arrays/_mixins.py
+++ b/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,
diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py
index 088c44334495c..34fa65a360b01 100644
--- a/pandas/core/arrays/base.py
+++ b/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,
diff --git a/pandas/core/arrays/string_arrow.py b/pandas/core/arrays/string_arrow.py
index 9411d3535e06f..4c8cb89378ecd 100644
--- a/pandas/core/arrays/string_arrow.py
+++ b/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,7 +511,10 @@ 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.
diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py
index 6887b919cc7d6..15cec99aff0eb 100644
--- a/pandas/core/indexes/base.py
+++ b/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)


Code review comments:
diff_hunk: @@ -216,6 +216,7 @@
 # Using List[int] here rather than Sequence[int] to disallow tuples.
 ScalarIndexer = Union[int, np.integer]
 SequenceIndexer = Union[slice, List[int], np.ndarray]
+TakeIndexer = Union[Sequence[int], np.ndarray]
path: pandas/_typing.py
text: @author:
Thank you! I committed your change and rebased.

@author:
Probably would be better to use:
```py
TakeIndexer = Union[Sequence[int], Sequence[np.integer], npt.NDArray[np.integer]]
```
to not allow a sequence of mixed types.