Closed
Description
The following code executes without any runtime error, but static type checkers such as mypy flag it as invalid:
import numpy as np
import pandas as pd
pd.DataFrame({"a": [1, 2, 3]}).take(np.array([0, 1]))
In pandas-stubs, the signature for DataFrame.take
currently restricts indices to list
:
def take(self, indices: list, axis: Axis = ..., **kwargs) -> Self: ...
However, the actual pandas implementation clearly accepts a NumPy array (and other array-like objects). For reference, the stub for _iLocIndexerFrame
already accommodates NumPy arrays.
class _iLocIndexerFrame(_iLocIndexer, Generic[_T]):
@overload
def __getitem__(
self,
idx: (
IndexType
| MaskType
| tuple[IndexType | MaskType, IndexType | MaskType]
| tuple[slice]
),
) -> _T: ...
It would be helpful if the take method signature could be fixed!
Metadata
Metadata
Assignees
Labels
No labels