ENH: Improve Arrays

pandas-stubs/_libs/interval.pyi

@@ -62,8 +62,6 @@ class IntervalMixin:
     def open_left(self) -> bool: ...
     @property
     def open_right(self) -> bool: ...
-    @property
-    def is_empty(self) -> bool: ...
 
 class Interval(IntervalMixin, Generic[_OrderableT]):
     @property
@@ -214,6 +212,8 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     def __ne__(self, other: Series[_OrderableT]) -> Series[bool]: ...  # type: ignore[misc]
     @overload
     def __ne__(self, other: object) -> Literal[True]: ...
+    @property
+    def is_empty(self) -> bool: ...
 
 class IntervalTree(IntervalMixin):
     def __init__(
@@ -233,3 +233,5 @@ class IntervalTree(IntervalMixin):
     @property
     def is_monotonic_increasing(self) -> bool: ...
     def clear_mapping(self) -> None: ...
+    @property
+    def is_empty(self) -> bool: ...

pandas-stubs/core/arrays/base.pyi

@@ -10,11 +10,11 @@ from pandas._typing import (
     npt,
 )
 
-from pandas.core.dtypes.dtypes import ExtensionDtype as ExtensionDtype
+from pandas.core.dtypes.dtypes import ExtensionDtype
 
 class ExtensionArray:
     def __getitem__(self, item) -> Any: ...
-    def __setitem__(self, key: int | slice | np.ndarray, value) -> None: ...
+    def __setitem__(self, key: int | slice | np.ndarray, value: Self) -> None: ...
     def __len__(self) -> int: ...
     def __iter__(self): ...
     def __contains__(self, item: object) -> bool | np.bool_: ...

pandas-stubs/core/arrays/boolean.pyi

@@ -1,28 +1,57 @@
-from typing import ClassVar
+from typing import (
+    ClassVar,
+    Sequence,
+    overload,
+)
 
 import numpy as np
-from pandas.core.arrays.masked import BaseMaskedArray as BaseMaskedArray
+from pandas.core.arrays import ExtensionArray
+from pandas.core.arrays.masked import BaseMaskedArray
+from typing_extensions import TypeAlias
 
 from pandas._libs.missing import NAType
-from pandas._typing import type_t
+from pandas._typing import (
+    npt,
+    type_t,
+)
 
-from pandas.core.dtypes.base import ExtensionDtype as ExtensionDtype
+from pandas.core.dtypes.base import ExtensionDtype
 
 class BooleanDtype(ExtensionDtype):
     na_value: ClassVar[NAType]
     @classmethod
     def construct_array_type(cls) -> type_t[BooleanArray]: ...
 
-def coerce_to_array(values, mask=..., copy: bool = ...): ...
+_ScalarType: TypeAlias = bool | np.bool_ | NAType | None
+_ArrayKey: TypeAlias = Sequence[int] | npt.NDArray[np.integer] | slice
 
 class BooleanArray(BaseMaskedArray):
     def __init__(
-        self, values: np.ndarray, mask: np.ndarray, copy: bool = ...
+        self,
+        values: npt.NDArray[np.bool_],
+        mask: npt.NDArray[np.bool_],
+        copy: bool = ...,
     ) -> None: ...
+    # Ignore overrides since more specific than super type
+    @overload  # type: ignore[override]
+    def __setitem__(self, key: int, value: _ScalarType) -> None: ...
+    @overload
+    def __setitem__(
+        self,
+        key: _ArrayKey,
+        value: _ScalarType
+        | Sequence[bool | NAType | None]
+        | npt.NDArray[np.bool_]
+        | BooleanArray,
+    ) -> None: ...
+    @overload
+    def __getitem__(self, item: int) -> bool | NAType: ...
+    @overload
+    def __getitem__(self, item: _ArrayKey) -> BooleanArray: ...
     @property
-    def dtype(self): ...
-    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): ...
-    def __setitem__(self, key, value) -> None: ...
-    def astype(self, dtype, copy: bool = ...): ...
-    def any(self, *, skipna: bool = ..., **kwargs): ...
-    def all(self, *, skipna: bool = ..., **kwargs): ...
+    def dtype(self) -> ExtensionDtype: ...
+    def astype(
+        self, dtype: npt.DTypeLike | ExtensionDtype, copy: bool = ...
+    ) -> np.ndarray | ExtensionArray: ...
+    def any(self, skipna: bool = ..., **kwargs) -> bool: ...
+    def all(self, skipna: bool = ..., **kwargs) -> bool: ...

pandas-stubs/core/arrays/integer.pyi

@@ -1,6 +1,20 @@
+from typing import (
+    Literal,
+    Sequence,
+    overload,
+)
+
+import numpy as np
+import pandas as pd
+from pandas.arrays import (
+    BooleanArray,
+    DatetimeArray,
+    StringArray,
+)
 from pandas.core.arrays.masked import BaseMaskedArray
 
 from pandas._libs.missing import NAType
+from pandas._typing import npt
 
 from pandas.core.dtypes.base import ExtensionDtype as ExtensionDtype
 
@@ -12,13 +26,96 @@ class _IntegerDtype(ExtensionDtype):
     def itemsize(self) -> int: ...
     @classmethod
     def construct_array_type(cls) -> type[IntegerArray]: ...
+    def __from_arrow__(self, array): ...
 
 class IntegerArray(BaseMaskedArray):
     def dtype(self): ...
-    def __init__(self, values, mask, copy: bool = ...) -> None: ...
-    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): ...
-    def __setitem__(self, key, value) -> None: ...
-    def astype(self, dtype, copy: bool = ...): ...
+    def __init__(
+        self,
+        values: npt.NDArray[np.integer],
+        mask: npt.NDArray[np.bool_] | Sequence[bool],
+        copy: bool = ...,
+    ) -> None: ...
+    @overload  # type: ignore[override]
+    def __setitem__(self, key: int, value: float | NAType) -> None: ...
+    @overload
+    def __setitem__(
+        self,
+        key: Sequence[int] | slice | npt.NDArray[np.integer],
+        value: float | NAType | Sequence[float | NAType] | npt.NDArray[np.integer],
+    ) -> None: ...
+    @overload
+    def __getitem__(self, item: int) -> int | NAType: ...
+    @overload
+    def __getitem__(
+        self, item: slice | list[int] | npt.NDArray[np.integer]
+    ) -> IntegerArray: ...
+    # Note: the ignores are needed below due to types being subclasses,
+    # e.g., float32 and float64 or bool, int, float, complex
+    @overload
+    def astype(  # type: ignore[misc]
+        self, dtype: Literal["str"] | type[str] | np.str_
+    ) -> npt.NDArray[np.str_]: ...
+    @overload
+    def astype(  # type: ignore[misc]
+        self, dtype: type[bool | np.bool_] | Literal["bool"], copy: bool = ...
+    ) -> npt.NDArray[np.bool_]: ...
+    @overload
+    def astype(
+        self,
+        dtype: Literal["i1", "i2", "i4", "i8", "int8", "int16", "int32", "int64"]
+        | type[int | np.signedinteger],
+        copy: bool = ...,
+    ) -> npt.NDArray[np.signedinteger]: ...
+    @overload
+    def astype(
+        self,
+        dtype: Literal["u1", "u2", "u4", "u8", "uint8", "uint16", "uint32", "uint64"]
+        | type[np.unsignedinteger],
+        copy: bool = ...,
+    ) -> npt.NDArray[np.unsignedinteger]: ...
+    @overload
+    def astype(  # type: ignore[misc]
+        self, dtype: Literal["f4", "float32"] | type[np.float32]
+    ) -> npt.NDArray[np.float32]: ...
+    @overload
+    def astype(
+        self, dtype: Literal["float", "float64", "f8"] | type[np.float64 | float]
+    ) -> npt.NDArray[np.float64]: ...
+    @overload
+    def astype(  # type: ignore[misc]
+        self, dtype: Literal["complex64", "c8"] | type[np.complex64]
+    ) -> npt.NDArray[np.complex64]: ...
+    @overload
+    def astype(
+        self,
+        dtype: Literal["complex", "complex128", "c16"] | type[np.complex128 | complex],
+    ) -> npt.NDArray[np.complex128]: ...
+    @overload
+    def astype(self, dtype: Literal["boolean"] | pd.BooleanDtype) -> BooleanArray: ...
+    @overload
+    def astype(
+        self,
+        dtype: Literal[
+            "Int8", "Int16", "Int32", "Int64", "UInt8", "UInt16", "UInt32", "UInt64"
+        ]
+        | pd.Int8Dtype
+        | pd.Int16Dtype
+        | pd.Int32Dtype
+        | pd.Int64Dtype
+        | pd.UInt8Dtype
+        | pd.UInt16Dtype
+        | pd.UInt32Dtype
+        | pd.UInt64Dtype,
+    ) -> IntegerArray: ...
+    @overload
+    def astype(self, dtype: Literal["string"] | pd.StringDtype) -> StringArray: ...
+    @overload
+    def astype(
+        self, dtype: type[np.datetime64] | Literal["M8[ns]"]
+    ) -> npt.NDArray[np.datetime64]: ...
+    @overload
+    def astype(self, dtype: pd.DatetimeTZDtype) -> DatetimeArray: ...
 
 class Int8Dtype(_IntegerDtype): ...
 class Int16Dtype(_IntegerDtype): ...