TYP: ExtensionArray.take accept np.ndarray
State: closed

- [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

Labels:
name: ExtensionArray
description: Extending pandas with custom dtypes or arrays.
name: Typing
description: type annotations, mypy/pyright type checking

Comments:
Using the `SequenceIndexer` from #41258. This means that `Sequence[int]` is not any more allowed (only `List[int], slice, np.ndarray`). I think any sequence should work for `take`, as all of them end up in `np.asarray`. Happy to allow `Sequence[int]` for `take` (but it would also be great to be consistent across methods: not allowing any sequence). I favor consistency.

@Dr-Irv
in the base class the docstring for take is `indices : sequence of int`. This is our contract with 3rd party EA developers.
> in the base class the docstring for take is `indices : sequence of int`. This is our contract with 3rd party EA developers.

does `Sequence[int]` *include* a suitable `ndarray`, e.g. `ndarray[intp]` (prob no?) if so we should extend `.take` to do this as its quite natural and IMHO expected.
from https://pandas.pydata.org/pandas-docs/dev/development/contributing_docstring.html#parameter-types

> If the exact type is not relevant, but must be compatible with a NumPy array, array-like can be specified. If Any type that can be iterated is accepted, iterable can be used:
>
> array-like
> iterable

having different rules for the docstring parameter types and the type annotations is bound to create confusion but IIUC we want to keep the docstring rules for consistency with the ecosystem.

There is no definition of the term "sequence" in the docstring guide.

np.ndarray is not an instance of abc.Sequence or type compatible with the typing protocol.

so if we want to annotate using Sequence we need to Union with np.ndarray (and Series, Index, EA too in most cases)

see #28770 for related issue.

I think the reason "sequence" is used in the docstring type parameters section instead of array-like is because we don't accept zero or multi dimension arrays. see #41807 for discussion.

see also #32773, https://github.com/pandas-dev/pandas/issues/32773#issuecomment-600064248, https://github.com/pandas-dev/pandas/issues/32773#issuecomment-840216855 and #32033 for discussions relating to avoiding the numpydoc type definitions for internal docstrings to therefore avoid the discrepancies (not relevant here as the base class docstrings are our api definitions for 3rd party EA developers)

if we could sometime in the future use the standard type notation in the public dosctrings that would be great (and redundant)
> I think the reason "sequence" is used in the docstring type parameters section instead of array-like is because we don't accept zero or multi dimension arrays. see #41807 for discussion.

Think we have two options then:

1.  Use `Sequence[int] | np.ndarray` and in the future add shape information to the `np.ndarray` when that becomes standardized in numpy
2.  Keep the current state (close this PR)
no need to close. accepting np.ndarray for take as Jeff says is sort of a given since `take` is effectively a numpy function.

numpy.take accepts array-like which can be any dimension. It only takes in 1 dimension but will populate a multidimensional array if the indices is mutli-dimensional. EAs a predominantly 1d (causes issues with \_\_array_function__ protocol extending the EA interface to be 2d is another topic)

The goal is to add types such that we don't add code in the main codebase that does not pass mypy silently when we break the EA contract with 3rd party EA developers (we effectively broke the contract in #41258 by adding `list[int]: A list of int` to the accepted types for \_\_getitem__, 3rd party code could break)

so here okay to add np.ndarray, but not okay to add slice, 3rd party EAs aren't expecting a slice. If we add to the type annotations and then pass a slice from the main code, mypy won't report any issues.
In general, type annotations should be as permissive as possible. we can't do that when typing the EA methods as we can't allow wider types than the 3rd party EA developers are expecting.
> I think the reason "sequence" is used in the docstring type parameters section instead of array-like is because we don't accept zero or multi dimension arrays. see #41807 for discussion.

I would question whether `take` should accept `Sequence[int]` or just `List[int]`, because I believe that `Sequence[int]` would include a tuple of ints, and if we want to be consistent between what `.iloc` and `.loc` do, those don't accept tuples.

So maybe `take` should only accept in its signature `List[int] | np.ndarray[int]` ??

> So maybe `take` should only accept in its signature `List[int] | np.ndarray[int]` ??

narrowing the types would not break 3rd party EAs but could break user code. We're quite constrained here without changing the EA api.

to be fair though EA.take is not part of the public api for users.

Index.take and Series.take could continue to be permissive and the user passed arg could be cast before dispatching to the EA.
> narrowing the types would not break 3rd party EAs but could break user code. We're quite constrained here without changing the EA api.

I think the question is what defines the EA API?  Is it the signature in the code?  Is it what is in the docs as the signature?  Is it what is in the docs that is in the docstring?

If you go to https://pandas.pydata.org/docs/reference/api/pandas.api.extensions.ExtensionArray.take.html?highlight=take#pandas.api.extensions.ExtensionArray.take

the docs don't define the types for the argument `indices`, even though the signature in the source says `Sequence[int]` .

If the signature in the code defines the API, are we expecting people to read those signatures by looking directly at pandas source?  Shouldn't the docs be the definition of the API?

Finally, I don't think changing the signature would break user code - it might just break it if they did type checking on their code.

> I think the question is what defines the EA API? Is it the signature in the code? Is it what is in the docs as the signature? Is it what is in the docs that is in the docstring?

none of the above. only jesting but probably the EA base tests should be in this list.

>If you go to https://pandas.pydata.org/docs/reference/api/pandas.api.extensions.ExtensionArray.take.html?highlight=take#pandas.api.extensions.ExtensionArray.take
>
>the docs don't define the types for the argument indices

it's defined in the parameter types section of the docstring. Type annotations across the codebase are incomplete and are not yet part of the official api. (IIRC we had a discussion in the past about removing the type annotations from the published docs)

> Finally, I don't think changing the signature would break user code - it might just break it if they did type checking on their code.

sure. but once we change the type annotations that's what mypy checks the code against for consistency. The point of inline type annotations is to prevent type related errors being introduced and latent bugs to be discovered. Adding type annotations as a sort of second level of documentation of what we would like the types to be renders the typing exercise pointless IMHO.

So if we change the annotations it won't immediately break user code (if no code changes are required to get mypy to green). given. but subsequent code changes could and we won't have mypy to stop us.



I think, except for maybe finding a better name for `TakeIndexer = Union[Sequence[int], np.ndarray]`, this PR should be ready. I also checked all the affected `take` docs strings, they all define the type as `sequence of int`.