Type ExtensionArray.{dropna,unique,repeat,take,ravel} (#1423)

* Type ``ExtensionArray.{dropna,unique,repeat,take,ravel}``

* introduce and use AnyArrayLikeInt

* add test file

* reorder

* reorder

* remove redundant

* remove redundant

* include Sequence, add comment about integerarray, reduce diff

Author: MarcoGorelli

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]: ...

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

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)