feat: complete ExtensionArrays

pandas-stubs/_libs/arrays.pyi

@@ -0,0 +1,38 @@
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+from typing_extensions import Self
+
+from pandas._typing import (
+    AnyArrayLikeInt,
+    AxisInt,
+    DtypeObj,
+    Shape,
+)
+
+class NDArrayBacked:
+    _dtype: DtypeObj
+    _ndarray: np.ndarray
+    def __setstate__(self, state: Any) -> None: ...
+    def __len__(self) -> int: ...
+    @property
+    def shape(self) -> Shape: ...
+    @property
+    def ndim(self) -> int: ...
+    @property
+    def size(self) -> int: ...
+    @property
+    def nbytes(self) -> int: ...
+    def copy(self, order=...) -> Self: ...
+    def delete(self, loc, axis=...) -> Self: ...
+    def swapaxes(self, axis1, axis2) -> Self: ...
+    def repeat(
+        self,
+        repeats: int | Sequence[int] | AnyArrayLikeInt,
+        axis: AxisInt | None = None,
+    ) -> Self: ...
+    def reshape(self, *args: Any, **kwargs: Any) -> Self: ...
+    def ravel(self, order=...) -> Self: ...
+    @property
+    def T(self) -> Self: ...

pandas-stubs/arrays/__init__.pyi

@@ -1,12 +1,31 @@
 from pandas.core.arrays import (
-    BooleanArray as BooleanArray,
-    Categorical as Categorical,
-    DatetimeArray as DatetimeArray,
-    IntegerArray as IntegerArray,
-    IntervalArray as IntervalArray,
-    PandasArray as PandasArray,
-    PeriodArray as PeriodArray,
-    SparseArray as SparseArray,
-    StringArray as StringArray,
-    TimedeltaArray as TimedeltaArray,
+    ArrowExtensionArray,
+    ArrowStringArray,
+    BooleanArray,
+    Categorical,
+    DatetimeArray,
+    FloatingArray,
+    IntegerArray,
+    IntervalArray,
+    NumpyExtensionArray,
+    PeriodArray,
+    SparseArray,
+    StringArray,
+    TimedeltaArray,
 )
+
+__all__ = [
+    "ArrowExtensionArray",
+    "ArrowStringArray",
+    "BooleanArray",
+    "Categorical",
+    "DatetimeArray",
+    "FloatingArray",
+    "IntegerArray",
+    "IntervalArray",
+    "NumpyExtensionArray",
+    "PeriodArray",
+    "SparseArray",
+    "StringArray",
+    "TimedeltaArray",
+]

pandas-stubs/core/algorithms.pyi

@@ -38,18 +38,19 @@ def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ...
 @overload
 def unique(values: PeriodIndex) -> PeriodIndex: ...
 @overload
-# switch to DatetimeIndex after Pandas 3.0
+# TODO: switch to DatetimeIndex after Pandas 3.0 pandas-dev/pandas#57064
 def unique(values: DatetimeIndex) -> np_1darray_dt | DatetimeIndex: ...
 @overload
-# switch to TimedeltaIndex after Pandas 3.0
+# TODO: switch to TimedeltaIndex after Pandas 3.0 pandas-dev/pandas#57064
 def unique(values: TimedeltaIndex) -> np_1darray_td: ...
 @overload
 # switch to Index[int] after Pandas 3.0
 def unique(values: RangeIndex) -> np_1darray_int64: ...
 @overload
 def unique(values: MultiIndex) -> np_ndarray: ...
 @overload
-def unique(values: Index) -> np_1darray | Index: ...  # switch to Index after Pandas 3.0
+# TODO: switch to Index after Pandas 3.0 pandas-dev/pandas#57064
+def unique(values: Index) -> np_1darray | Index: ...
 @overload
 def unique(values: Categorical) -> Categorical: ...
 

pandas-stubs/core/arrays/__init__.pyi

@@ -1,15 +1,39 @@
+from pandas.core.arrays.arrow import ArrowExtensionArray
 from pandas.core.arrays.base import (
-    ExtensionArray as ExtensionArray,
-    ExtensionOpsMixin as ExtensionOpsMixin,
-    ExtensionScalarOpsMixin as ExtensionScalarOpsMixin,
+    ExtensionArray,
+    ExtensionOpsMixin,
+    ExtensionScalarOpsMixin,
 )
-from pandas.core.arrays.boolean import BooleanArray as BooleanArray
-from pandas.core.arrays.categorical import Categorical as Categorical
-from pandas.core.arrays.datetimes import DatetimeArray as DatetimeArray
-from pandas.core.arrays.integer import IntegerArray as IntegerArray
-from pandas.core.arrays.interval import IntervalArray as IntervalArray
-from pandas.core.arrays.numpy_ import PandasArray as PandasArray
-from pandas.core.arrays.period import PeriodArray as PeriodArray
-from pandas.core.arrays.sparse import SparseArray as SparseArray
-from pandas.core.arrays.string_ import StringArray as StringArray
-from pandas.core.arrays.timedeltas import TimedeltaArray as TimedeltaArray
+from pandas.core.arrays.boolean import BooleanArray
+from pandas.core.arrays.categorical import Categorical
+from pandas.core.arrays.datetimes import DatetimeArray
+from pandas.core.arrays.floating import FloatingArray
+from pandas.core.arrays.integer import IntegerArray
+from pandas.core.arrays.interval import IntervalArray
+from pandas.core.arrays.masked import BaseMaskedArray
+from pandas.core.arrays.numpy_ import NumpyExtensionArray
+from pandas.core.arrays.period import PeriodArray
+from pandas.core.arrays.sparse import SparseArray
+from pandas.core.arrays.string_ import StringArray
+from pandas.core.arrays.string_arrow import ArrowStringArray
+from pandas.core.arrays.timedeltas import TimedeltaArray
+
+__all__ = [
+    "ArrowExtensionArray",
+    "ArrowStringArray",
+    "BaseMaskedArray",
+    "BooleanArray",
+    "Categorical",
+    "DatetimeArray",
+    "ExtensionArray",
+    "ExtensionOpsMixin",
+    "ExtensionScalarOpsMixin",
+    "FloatingArray",
+    "IntegerArray",
+    "IntervalArray",
+    "NumpyExtensionArray",
+    "PeriodArray",
+    "SparseArray",
+    "StringArray",
+    "TimedeltaArray",
+]

pandas-stubs/core/arrays/_arrow_string_mixins.pyi

@@ -0,0 +1 @@
+class ArrowStringArrayMixin: ...

pandas-stubs/core/arrays/_mixins.pyi

@@ -0,0 +1,15 @@
+from pandas.core.arrays.base import ExtensionArray
+from pandas.core.series import Series
+from typing_extensions import Self
+
+from pandas._libs.arrays import NDArrayBacked
+from pandas._typing import (
+    AxisInt,
+    Scalar,
+)
+
+class NDArrayBackedExtensionArray(NDArrayBacked, ExtensionArray):
+    def argmin(self, axis: AxisInt = 0, skipna: bool = True) -> int: ...
+    def argmax(self, axis: AxisInt = 0, skipna: bool = True) -> int: ...
+    def insert(self, loc: int, item: Scalar) -> Self: ...
+    def value_counts(self, dropna: bool = True) -> Series[int]: ...

pandas-stubs/core/arrays/arrow/__init__.pyi

@@ -0,0 +1,7 @@
+from pandas.core.arrays.arrow.accessors import (
+    ListAccessor,
+    StructAccessor,
+)
+from pandas.core.arrays.arrow.array import ArrowExtensionArray
+
+__all__ = ["ArrowExtensionArray", "ListAccessor", "StructAccessor"]

pandas-stubs/core/arrays/arrow/accessors.pyi

@@ -0,0 +1,5 @@
+from abc import ABCMeta
+
+class ArrowAccessor(metaclass=ABCMeta): ...
+class ListAccessor(ArrowAccessor): ...
+class StructAccessor(ArrowAccessor): ...

pandas-stubs/core/arrays/arrow/array.pyi

@@ -0,0 +1,7 @@
+from pandas.core.arraylike import OpsMixin
+from pandas.core.arrays._arrow_string_mixins import ArrowStringArrayMixin
+from pandas.core.arrays.base import ExtensionArraySupportsAnyAll
+
+class ArrowExtensionArray(
+    OpsMixin, ExtensionArraySupportsAnyAll, ArrowStringArrayMixin
+): ...

pandas-stubs/core/arrays/base.pyi

@@ -87,7 +87,7 @@ class ExtensionArray:
     ) -> np.intp: ...
     def factorize(self, use_na_sentinel: bool = True) -> tuple[np_1darray, Self]: ...
     def repeat(
-        self, repeats: int | AnyArrayLikeInt | Sequence[int], axis: None = None
+        self, repeats: int | Sequence[int] | AnyArrayLikeInt, axis: None = None
     ) -> Self: ...
     def take(
         self,
@@ -114,6 +114,8 @@ class ExtensionArray:
         **kwargs: Any,
     ) -> Self: ...
 
+class ExtensionArraySupportsAnyAll(ExtensionArray): ...
+
 class ExtensionOpsMixin:
     @classmethod
     def _add_arithmetic_ops(cls) -> None: ...

pandas-stubs/core/arrays/boolean.pyi

@@ -1,9 +1,10 @@
 from collections.abc import Sequence
-from typing import Any
 
 import numpy as np
 from pandas.core.arrays.integer import IntegerArray
 from pandas.core.arrays.masked import BaseMaskedArray as BaseMaskedArray
+from pandas.core.indexes.base import Index
+from pandas.core.series import Series
 from typing_extensions import Self
 
 from pandas._libs.missing import NAType
@@ -22,13 +23,21 @@ class BooleanDtype(ExtensionDtype):
 
 class BooleanArray(BaseMaskedArray):
     def __init__(
-        self, values: np_ndarray_bool, mask: np_ndarray_bool, copy: bool = ...
+        self,
+        values: (
+            Sequence[bool | np.bool]
+            | np_ndarray_bool
+            | Index[bool]
+            | Series[bool]
+            | Self
+        ),
+        mask: np_ndarray_bool,
+        copy: bool = False,
     ) -> None: ...
     @property
-    def dtype(self): ...
-    def __setitem__(self, key, value) -> None: ...
-    def any(self, *, skipna: bool = ..., **kwargs: Any): ...
-    def all(self, *, skipna: bool = ..., **kwargs: Any): ...
+    def dtype(  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+    ) -> np.dtypes.BoolDType: ...
     def __and__(
         self,
         other: (

pandas-stubs/core/arrays/datetimes.pyi

@@ -1,73 +1,33 @@
-from datetime import tzinfo as _tzinfo
+from collections.abc import Sequence
+from datetime import datetime
 
 import numpy as np
 from pandas.core.arrays.datetimelike import (
     DatelikeOps,
     DatetimeLikeArrayMixin,
     TimelikeOps,
 )
+from pandas.core.indexes.datetimes import DatetimeIndex
+from pandas.core.series import Series
+from typing_extensions import Self
 
-from pandas._typing import (
-    TimeAmbiguous,
-    TimeNonexistent,
-    TimeZones,
-)
-
-from pandas.core.dtypes.dtypes import DatetimeTZDtype as DatetimeTZDtype
+from pandas._libs.tslibs.timestamps import Timestamp
 
 class DatetimeArray(DatetimeLikeArrayMixin, TimelikeOps, DatelikeOps):
     __array_priority__: int = ...
-    def __init__(self, values, dtype=..., freq=..., copy: bool = ...) -> None: ...
-    # ignore in dtype() is from the pandas source
-    @property
-    def dtype(self) -> np.dtype | DatetimeTZDtype: ...  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
-    @property
-    def tz(self): ...
-    @tz.setter
-    def tz(self, value) -> None: ...
-    @property
-    def tzinfo(self) -> _tzinfo | None: ...
-    @property
-    def is_normalized(self): ...
-    def __iter__(self): ...
-    def tz_convert(self, tz: TimeZones): ...
-    def tz_localize(
+    def __init__(
         self,
-        tz: TimeZones,
-        ambiguous: TimeAmbiguous = "raise",
-        nonexistent: TimeNonexistent = "raise",
-    ): ...
-    def to_pydatetime(self): ...
-    def normalize(self): ...
-    def to_period(self, freq=...): ...
-    def to_perioddelta(self, freq): ...
-    def month_name(self, locale=...): ...
-    def day_name(self, locale=...): ...
-    @property
-    def time(self): ...
+        values: (
+            Sequence[datetime | np.datetime64]
+            | np.typing.NDArray[np.datetime64]
+            | DatetimeIndex
+            | Series[Timestamp]
+            | Self
+        ),
+        dtype: np.dtype | None = None,
+        copy: bool = False,
+    ) -> None: ...
     @property
-    def timetz(self): ...
-    @property
-    def date(self): ...
-    year = ...
-    month = ...
-    day = ...
-    hour = ...
-    minute = ...
-    second = ...
-    microsecond = ...
-    nanosecond = ...
-    dayofweek = ...
-    weekday = ...
-    dayofyear = ...
-    quarter = ...
-    days_in_month = ...
-    daysinmonth = ...
-    is_month_start = ...
-    is_month_end = ...
-    is_quarter_start = ...
-    is_quarter_end = ...
-    is_year_start = ...
-    is_year_end = ...
-    is_leap_year = ...
-    def to_julian_date(self): ...
+    def dtype(  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+    ) -> np.dtypes.DateTime64DType: ...

pandas-stubs/core/arrays/integer.pyi

@@ -1,3 +1,4 @@
+import numpy as np
 from pandas.core.arrays.masked import BaseMaskedArray
 
 from pandas._libs.missing import NAType
@@ -16,8 +17,12 @@ class _IntegerDtype(ExtensionDtype):
 class IntegerArray(BaseMaskedArray):
     @property
     def dtype(self) -> _IntegerDtype: ...
-    def __init__(self, values, mask, copy: bool = ...) -> None: ...
-    def __setitem__(self, key, value) -> None: ...
+    def __init__(
+        self,
+        values: np.typing.NDArray[np.integer],
+        mask: np.typing.NDArray[np.bool],
+        copy: bool = False,
+    ) -> None: ...
 
 class Int8Dtype(_IntegerDtype): ...
 class Int16Dtype(_IntegerDtype): ...

pandas-stubs/core/arrays/numpy_.pyi

@@ -1,16 +1,7 @@
-import numpy as np
-from numpy.lib.mixins import NDArrayOperatorsMixin
-from pandas.core.arrays.base import (
-    ExtensionArray,
-    ExtensionOpsMixin,
-)
+from pandas.core.arraylike import OpsMixin
+from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
+from pandas.core.strings.object_array import ObjectStringArrayMixin
 
-from pandas.core.dtypes.dtypes import ExtensionDtype
-
-class PandasDtype(ExtensionDtype):
-    @property
-    def numpy_dtype(self) -> np.dtype: ...
-    @property
-    def itemsize(self) -> int: ...
-
-class PandasArray(ExtensionArray, ExtensionOpsMixin, NDArrayOperatorsMixin): ...
+class NumpyExtensionArray(
+    OpsMixin, NDArrayBackedExtensionArray, ObjectStringArrayMixin
+): ...