Skip to content

pd.DataFrame.take should allow np.ndarray indices (current stub only accepts list) #1208

Closed
@ritsuki1227

Description

@ritsuki1227

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

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions