diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 2c72f8d56..61204d898 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -28,7 +28,10 @@ from typing import ( import numpy as np from numpy import typing as npt import pandas as pd -from pandas.core.arrays import ExtensionArray +from pandas.core.arrays import ( + ExtensionArray, + IntegerArray, +) from pandas.core.frame import DataFrame from pandas.core.generic import NDFrame from pandas.core.groupby.grouper import Grouper @@ -861,6 +864,10 @@ np_ndarray: TypeAlias = np.ndarray[ShapeT, np.dtype[GenericT]] np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[GenericT]] np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]] +AnyArrayLikeInt: TypeAlias = ( + IntegerArray | Index[int] | Series[int] | np_1darray[np.integer] | Sequence[int] +) + class SupportsDType(Protocol[GenericT_co]): @property def dtype(self) -> np.dtype[GenericT_co]: ... diff --git a/pandas-stubs/core/arrays/base.pyi b/pandas-stubs/core/arrays/base.pyi index 6b3a48226..3f9de529e 100644 --- a/pandas-stubs/core/arrays/base.pyi +++ b/pandas-stubs/core/arrays/base.pyi @@ -9,6 +9,7 @@ import numpy as np from typing_extensions import Self from pandas._typing import ( + AnyArrayLikeInt, ArrayLike, Scalar, ScalarIndexer, @@ -52,22 +53,22 @@ class ExtensionArray: self, *, ascending: bool = ..., kind: str = ..., **kwargs: Any ) -> np_1darray: ... def fillna(self, value=..., method=None, limit=None): ... - def dropna(self): ... + def dropna(self) -> Self: ... def shift(self, periods: int = 1, fill_value: object = ...) -> Self: ... - def unique(self): ... + def unique(self) -> Self: ... def searchsorted(self, value, side: str = ..., sorter=...): ... def factorize(self, use_na_sentinel: bool = True) -> tuple[np_1darray, Self]: ... - def repeat(self, repeats, axis=...): ... + def repeat(self, repeats: int | AnyArrayLikeInt, axis: None = None) -> Self: ... def take( self, indexer: TakeIndexer, *, - allow_fill: bool = ..., - fill_value=..., + allow_fill: bool = False, + fill_value: Any = None, ) -> Self: ... def copy(self) -> Self: ... def view(self, dtype=...) -> Self | np_1darray: ... - def ravel(self, order="C") -> Self: ... + def ravel(self, order: Literal["C", "F", "A", "K"] | None = "C") -> Self: ... def tolist(self) -> list: ... def _reduce( self, name: str, *, skipna: bool = ..., keepdims: bool = ..., **kwargs: Any diff --git a/tests/arrays/test_extension_array.py b/tests/arrays/test_extension_array.py new file mode 100644 index 000000000..ed5638b0e --- /dev/null +++ b/tests/arrays/test_extension_array.py @@ -0,0 +1,31 @@ +# Test common ExtensionArray methods + +import pandas as pd +from pandas.core.arrays.integer import IntegerArray +from pandas.core.construction import array +from typing_extensions import assert_type + +from tests import check + + +def test_ea_common() -> None: + # Note: `ExtensionArray` is abstract, so we use `IntegerArray` for the tests. + arr = array([1, 2, 3]) + + check(assert_type(arr.repeat(1), IntegerArray), IntegerArray) + check(assert_type(arr.repeat(arr), IntegerArray), IntegerArray) + check( + assert_type(arr.repeat(repeats=pd.Series([1, 2, 3])), IntegerArray), + IntegerArray, + ) + check(assert_type(arr.repeat(pd.Index([1, 2, 3])), IntegerArray), IntegerArray) + check(assert_type(arr.repeat([1, 2, 3]), IntegerArray), IntegerArray) + + check(assert_type(arr.unique(), IntegerArray), IntegerArray) + check(assert_type(arr.dropna(), IntegerArray), IntegerArray) + check(assert_type(arr.take([1, 0, 2]), IntegerArray), IntegerArray) + check( + assert_type(arr.take([1, 0, 2], allow_fill=True, fill_value=-1), IntegerArray), + IntegerArray, + ) + check(assert_type(arr.ravel(), IntegerArray), IntegerArray)