Support an ExtensionDtype and ExtensionArray

pandas-stubs/_libs/lib.pyi

@@ -1,9 +1,28 @@
+from enum import Enum
+from typing import (
+    Final,
+    Literal,
+)
+
+import numpy as np
+from pandas import Interval
+from typing_extensions import (
+    TypeAlias,
+    TypeGuard,
+)
+
+class _NoDefault(Enum):
+    no_default = ...
+
+no_default: Final = _NoDefault.no_default
+NoDefault: TypeAlias = Literal[_NoDefault.no_default]
+
 def infer_dtype(value: object, skipna: bool = ...) -> str: ...
 def is_iterator(obj: object) -> bool: ...
 def is_scalar(val: object) -> bool: ...
 def is_list_like(obj: object, allow_sets: bool = ...) -> bool: ...
-def is_interval(val: object) -> bool: ...
-def is_complex(val: object) -> bool: ...
-def is_bool(val: object) -> bool: ...
-def is_integer(val: object) -> bool: ...
-def is_float(val: object) -> bool: ...
+def is_interval(val: object) -> TypeGuard[Interval]: ...
+def is_complex(val: object) -> TypeGuard[complex]: ...
+def is_bool(val: object) -> TypeGuard[bool | np.bool_]: ...
+def is_integer(val: object) -> TypeGuard[int | np.integer]: ...
+def is_float(val: object) -> TypeGuard[float | np.floating]: ...

pandas-stubs/_typing.pyi

@@ -166,7 +166,8 @@ class WriteExcelBuffer(WriteBuffer[bytes], Protocol):
 
 FilePath: TypeAlias = str | PathLike[str]
 
-Axis: TypeAlias = str | int
+AxisInt: TypeAlias = Literal[0, 1]
+Axis: TypeAlias = AxisInt | Literal["index", "columns", "rows"]
 IndexLabel: TypeAlias = Hashable | Sequence[Hashable]
 Label: TypeAlias = Hashable | None
 Level: TypeAlias = Hashable | int
@@ -293,6 +294,8 @@ IntervalT = TypeVar(
 )
 IntervalClosedType: TypeAlias = Literal["left", "right", "both", "neither"]
 
+TakeIndexer: TypeAlias = Sequence[int] | Sequence[np.integer] | npt.NDArray[np.integer]
+
 IgnoreRaiseCoerce: TypeAlias = Literal["ignore", "raise", "coerce"]
 
 # Shared by functions such as drop and astype

pandas-stubs/api/extensions/__init__.pyi

@@ -3,11 +3,14 @@ from pandas.core.accessor import (
     register_index_accessor as register_index_accessor,
     register_series_accessor as register_series_accessor,
 )
+from pandas.core.algorithms import take as take
 from pandas.core.arrays import (
     ExtensionArray as ExtensionArray,
     ExtensionScalarOpsMixin as ExtensionScalarOpsMixin,
 )
 
+from pandas._libs.lib import no_default as no_default
+
 from pandas.core.dtypes.dtypes import (
     ExtensionDtype as ExtensionDtype,
     register_extension_dtype as register_extension_dtype,

pandas-stubs/core/algorithms.pyi

@@ -14,7 +14,9 @@ from pandas.api.extensions import ExtensionArray
 
 from pandas._typing import (
     AnyArrayLike,
+    AxisInt,
     IntervalT,
+    TakeIndexer,
 )
 
 # These are type: ignored because the Index types overlap due to inheritance but indices
@@ -69,3 +71,10 @@ def value_counts(
     bins: int | None = ...,
     dropna: bool = ...,
 ) -> Series: ...
+def take(
+    arr,
+    indices: TakeIndexer,
+    axis: AxisInt = 0,
+    allow_fill: bool = False,
+    fill_value=None,
+): ...

pandas-stubs/core/arraylike.pyi

@@ -2,6 +2,10 @@ from typing import Any
 
 from typing_extensions import Self
 
+from pandas._libs.ops_dispatch import (
+    maybe_dispatch_ufunc_to_dunder_op as maybe_dispatch_ufunc_to_dunder_op,
+)
+
 class OpsMixin:
     def __eq__(self, other: object) -> Self: ...  # type: ignore[override]
     def __ne__(self, other: object) -> Self: ...  # type: ignore[override]

pandas-stubs/core/arrays/base.pyi

@@ -1,21 +1,23 @@
-from collections.abc import Sequence
+from typing import Any
 
 import numpy as np
+from typing_extensions import Self
 
 from pandas._typing import (
     ArrayLike,
     Scalar,
+    TakeIndexer,
     npt,
 )
 
 from pandas.core.dtypes.dtypes import ExtensionDtype as ExtensionDtype
-from pandas.core.dtypes.generic import ABCExtensionArray
 
 class ExtensionArray:
-    def __getitem__(self, item) -> None: ...
+    def __getitem__(self, item) -> Any: ...
     def __setitem__(self, key: int | slice | np.ndarray, value) -> None: ...
     def __len__(self) -> int: ...
     def __iter__(self): ...
+    def __contains__(self, item: object) -> bool | np.bool_: ...
     def to_numpy(
         self,
         dtype: npt.DTypeLike | None = ...,
@@ -37,22 +39,29 @@ class ExtensionArray:
     ) -> np.ndarray: ...
     def fillna(self, value=..., method=..., limit=...): ...
     def dropna(self): ...
-    def shift(
-        self, periods: int = ..., fill_value: object = ...
-    ) -> ABCExtensionArray: ...
+    def shift(self, periods: int = ..., fill_value: object = ...) -> Self: ...
     def unique(self): ...
     def searchsorted(self, value, side: str = ..., sorter=...): ...
     # TODO: remove keyword-only when pandas removed na_sentinel
-    def factorize(
-        self, *, use_na_sentinel: bool = ...
-    ) -> tuple[np.ndarray, ABCExtensionArray]: ...
+    def factorize(self, *, use_na_sentinel: bool = ...) -> tuple[np.ndarray, Self]: ...
     def repeat(self, repeats, axis=...): ...
     def take(
-        self, indices: Sequence[int], *, allow_fill: bool = ..., fill_value=...
-    ) -> ABCExtensionArray: ...
-    def copy(self) -> ABCExtensionArray: ...
-    def view(self, dtype=...) -> ABCExtensionArray | np.ndarray: ...
-    def ravel(self, order=...) -> ABCExtensionArray: ...
+        self,
+        indexer: TakeIndexer,
+        *,
+        allow_fill: bool = ...,
+        fill_value=...,
+    ) -> Self: ...
+    def copy(self) -> Self: ...
+    def view(self, dtype=...) -> Self | np.ndarray: ...
+    def ravel(self, order=...) -> Self: ...
+
+class ExtensionOpsMixin:
+    @classmethod
+    def _add_arithmetic_ops(cls) -> None: ...
+    @classmethod
+    def _add_comparison_ops(cls) -> None: ...
+    @classmethod
+    def _add_logical_ops(cls) -> None: ...
 
-class ExtensionOpsMixin: ...
 class ExtensionScalarOpsMixin(ExtensionOpsMixin): ...

pandas-stubs/core/arrays/boolean.pyi

@@ -1,3 +1,5 @@
+from typing import ClassVar
+
 import numpy as np
 
 from pandas._libs.missing import NAType
@@ -8,9 +10,7 @@ from pandas.core.dtypes.base import ExtensionDtype as ExtensionDtype
 from .masked import BaseMaskedArray as BaseMaskedArray
 
 class BooleanDtype(ExtensionDtype):
-    name: str = ...
-    @property
-    def na_value(self) -> NAType: ...
+    na_value: ClassVar[NAType]
     @classmethod
     def construct_array_type(cls) -> type_t[BooleanArray]: ...
 

pandas-stubs/core/arrays/categorical.pyi

@@ -24,6 +24,7 @@ from pandas._typing import (
     ListLike,
     Ordered,
     Scalar,
+    TakeIndexer,
     np_ndarray_bool,
     np_ndarray_int,
 )
@@ -165,7 +166,9 @@ class Categorical(ExtensionArray, PandasObject):
     def view(self, dtype=...): ...
     def to_dense(self): ...
     def fillna(self, value=..., method=..., limit=...): ...
-    def take(self, indexer, *, allow_fill: bool = ..., fill_value=...): ...
+    def take(
+        self, indexer: TakeIndexer, *, allow_fill: bool = ..., fill_value=...
+    ) -> Categorical: ...
     def take_nd(self, indexer, allow_fill: bool = ..., fill_value=...): ...
     def __len__(self) -> int: ...
     def __iter__(self): ...

pandas-stubs/core/arrays/datetimelike.pyi

@@ -5,11 +5,13 @@ from pandas.core.arrays.base import (
     ExtensionArray,
     ExtensionOpsMixin,
 )
+from typing_extensions import Self
 
 from pandas._libs import (
     NaT as NaT,
     NaTType as NaTType,
 )
+from pandas._typing import TakeIndexer
 
 class DatelikeOps:
     def strftime(self, date_format): ...
@@ -40,7 +42,9 @@ class DatetimeLikeArrayMixin(ExtensionOpsMixin, ExtensionArray):
     def astype(self, dtype, copy: bool = ...): ...
     def view(self, dtype=...): ...
     def unique(self): ...
-    def take(self, indices, *, allow_fill: bool = ..., fill_value=...): ...
+    def take(
+        self: Self, indices: TakeIndexer, *, allow_fill: bool = ..., fill_value=...
+    ) -> Self: ...
     def copy(self): ...
     def shift(self, periods: int = ..., fill_value=..., axis: int = ...): ...
     def searchsorted(self, value, side: str = ..., sorter=...): ...

pandas-stubs/core/arrays/interval.pyi

@@ -1,14 +1,16 @@
 import numpy as np
 from pandas import Index
 from pandas.core.arrays.base import ExtensionArray as ExtensionArray
+from typing_extensions import Self
 
 from pandas._libs.interval import (
     Interval as Interval,
     IntervalMixin as IntervalMixin,
 )
-from pandas._typing import Axis
-
-from pandas.core.dtypes.generic import ABCExtensionArray
+from pandas._typing import (
+    Axis,
+    TakeIndexer,
+)
 
 class IntervalArray(IntervalMixin, ExtensionArray):
     ndim: int = ...
@@ -40,12 +42,16 @@ class IntervalArray(IntervalMixin, ExtensionArray):
     def nbytes(self) -> int: ...
     @property
     def size(self) -> int: ...
-    def shift(
-        self, periods: int = ..., fill_value: object = ...
-    ) -> ABCExtensionArray: ...
-    def take(
-        self, indices, *, allow_fill: bool = ..., fill_value=..., axis=..., **kwargs
-    ): ...
+    def shift(self, periods: int = ..., fill_value: object = ...) -> IntervalArray: ...
+    def take(  # type: ignore[override]
+        self: Self,
+        indices: TakeIndexer,
+        *,
+        allow_fill: bool = ...,
+        fill_value=...,
+        axis=...,
+        **kwargs,
+    ) -> Self: ...
     def value_counts(self, dropna: bool = ...): ...
     @property
     def left(self) -> Index: ...

pandas-stubs/core/arrays/masked.pyi

@@ -6,6 +6,7 @@ from pandas.core.arrays import (
 
 from pandas._typing import (
     Scalar,
+    TakeIndexer,
     npt,
 )
 
@@ -26,6 +27,8 @@ class BaseMaskedArray(ExtensionArray, ExtensionOpsMixin):
     def isna(self): ...
     @property
     def nbytes(self) -> int: ...
-    def take(self, indexer, *, allow_fill: bool = ..., fill_value=...): ...
+    def take(
+        self, indexer: TakeIndexer, allow_fill: bool = ..., fill_value=...
+    ) -> BaseMaskedArray: ...
     def copy(self): ...
     def value_counts(self, dropna: bool = ...): ...

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

@@ -5,6 +5,8 @@ from pandas.core.arrays import (
 )
 from pandas.core.base import PandasObject
 
+from pandas._typing import TakeIndexer
+
 class SparseArray(PandasObject, ExtensionArray, ExtensionOpsMixin):
     def __init__(
         self,
@@ -42,12 +44,14 @@ class SparseArray(PandasObject, ExtensionArray, ExtensionOpsMixin):
     def fillna(self, value=..., method=..., limit=...): ...
     def shift(self, periods: int = ..., fill_value=...): ...
     def unique(self): ...
-    def factorize(  # type: ignore[override]
+    def factorize(
         self, na_sentinel: int = ..., use_na_sentinel: bool = ...
     ) -> tuple[np.ndarray, SparseArray]: ...
     def value_counts(self, dropna: bool = ...): ...
     def __getitem__(self, key): ...
-    def take(self, indices, *, allow_fill: bool = ..., fill_value=...): ...
+    def take(
+        self, indices: TakeIndexer, *, allow_fill: bool = ..., fill_value=...
+    ) -> SparseArray: ...
     def searchsorted(self, v, side: str = ..., sorter=...): ...
     def copy(self): ...
     def astype(self, dtype=..., copy: bool = ...): ...

pandas-stubs/core/dtypes/base.pyi

@@ -1,23 +1,23 @@
-from typing import Literal
+from typing import (
+    ClassVar,
+    Literal,
+)
 
 from pandas.core.arrays import ExtensionArray
 
-from pandas._libs import NaTType
-from pandas._libs.missing import NAType
 from pandas._typing import type_t
 
 class ExtensionDtype:
+    type: ClassVar[type_t]
+    name: ClassVar[str]
+
     @property
-    def na_value(self) -> NAType | NaTType: ...
-    @property
-    def type(self) -> type_t: ...
+    def na_value(self) -> object: ...
     @property
     def kind(
         self,
     ) -> Literal["b", "i", "u", "f", "c", "m", "M", "O", "S", "U", "V"]: ...
     @property
-    def name(self) -> str: ...
-    @property
     def names(self) -> list[str] | None: ...
     def empty(self, size: int | tuple[int, ...]) -> type_t[ExtensionArray]: ...
     @classmethod

pandas-stubs/core/dtypes/dtypes.pyi

@@ -2,6 +2,7 @@ import datetime as dt
 from typing import (
     Any,
     Literal,
+    TypeVar,
 )
 
 import numpy as np
@@ -17,7 +18,9 @@ from pandas._typing import (
 
 from .base import ExtensionDtype as ExtensionDtype
 
-def register_extension_dtype(cls: type[ExtensionDtype]) -> type[ExtensionDtype]: ...
+_ExtensionDtypeT = TypeVar("_ExtensionDtypeT", bound=ExtensionDtype)
+
+def register_extension_dtype(cls: type[_ExtensionDtypeT]) -> type[_ExtensionDtypeT]: ...
 
 class BaseMaskedDtype(ExtensionDtype): ...
 class PandasExtensionDtype(ExtensionDtype): ...

pyproject.toml

@@ -38,7 +38,7 @@ types-pytz = ">= 2022.1.1"
 mypy = "1.0"
 pyarrow = ">=10.0.1"
 pytest = ">=7.1.2"
-pyright = ">=1.1.286"
+pyright = ">=1.1.295"
 poethepoet = ">=0.16.5"
 loguru = ">=0.6.0"
 pandas = "1.5.3"