ENH: Improve Arrays

Author: bashtage

pandas-stubs/_libs/tslibs/period.pyi

@@ -2,11 +2,13 @@ from typing import Any
 
 class IncompatibleFrequency(ValueError): ...
 
+from pandas._libs.tslibs.offsets import BaseOffset
+
 class Period:
     def __init__(
         self,
         value: Any = ...,
-        freqstr: Any = ...,
+        freq: str | BaseOffset | None = ...,
         ordinal: Any = ...,
         year: Any = ...,
         month: int = ...,

pandas-stubs/core/arrays/boolean.pyi

@@ -1,4 +1,5 @@
 import numpy as np
+from pandas.core.arrays import ExtensionArray
 
 from pandas._typing import (
     Scalar,
@@ -21,16 +22,15 @@ class BooleanDtype(ExtensionDtype):
     def construct_array_type(cls) -> type_t[BooleanArray]: ...
     def __from_arrow__(self, array): ...
 
-def coerce_to_array(values, mask=..., copy: bool = ...): ...
-
 class BooleanArray(BaseMaskedArray):
     def __init__(
         self, values: np.ndarray, mask: np.ndarray, copy: bool = ...
     ) -> None: ...
+    def __setitem__(self, key: int | np.ndarray | slice, value: object) -> None: ...
     @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 astype(
+        self, dtype: str | np.dtype, 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,3 +1,23 @@
+from typing import (
+    Literal,
+    Sequence,
+    overload,
+)
+
+import numpy as np
+import pandas as pd
+from pandas.arrays import (
+    BooleanArray,
+    DatetimeArray,
+    StringArray,
+)
+
+from pandas._libs.missing import NAType
+from pandas._typing import (
+    Dtype,
+    npt,
+)
+
 from pandas.core.dtypes.base import ExtensionDtype as ExtensionDtype
 
 from .masked import BaseMaskedArray
@@ -18,15 +38,103 @@ class _IntegerDtype(ExtensionDtype):
     def construct_array_type(cls): ...
     def __from_arrow__(self, array): ...
 
-def safe_cast(values, dtype, copy): ...
-def coerce_to_array(values, dtype, mask=..., copy: bool = ...): ...
-
 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  # type: ignore[override]
+    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] | Literal["bool"] | type[np.bool_], copy: bool = ...
+    ) -> npt.NDArray[np.bool_]: ...
+    @overload
+    def astype(
+        self,
+        dtype: type[int]
+        | Literal["i1", "i2", "i4", "i8", "int8", "int16", "int32", "int64"]
+        | type[np.int8]
+        | type[np.int16]
+        | type[np.int32]
+        | type[np.int64],
+        copy: bool = ...,
+    ) -> npt.NDArray[np.signedinteger]: ...
+    @overload
+    def astype(
+        self,
+        dtype: Literal["u1", "u2", "u4", "u8", "uint8", "uint16", "uint32", "uint64"]
+        | type[np.uint8]
+        | type[np.uint16]
+        | type[np.uint32]
+        | type[np.uint64],
+        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: type[float] | Literal["float", "float64", "f8"] | type[np.float64]
+    ) -> 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: type[complex]
+        | Literal["complex", "complex128", "c16"]
+        | type[np.complex128],
+    ) -> 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): ...

pandas-stubs/core/arrays/interval.pyi

@@ -17,13 +17,17 @@ class IntervalArray(IntervalMixin, ExtensionArray):
         cls, data, closed=..., dtype=..., copy: bool = ..., verify_integrity: bool = ...
     ): ...
     @classmethod
-    def from_breaks(cls, breaks, closed: str = ..., copy: bool = ..., dtype=...): ...
+    def from_breaks(
+        cls, breaks, closed: str = ..., copy: bool = ..., dtype=...
+    ) -> IntervalArray: ...
     @classmethod
     def from_arrays(
         cls, left, right, closed: str = ..., copy: bool = ..., dtype=...
-    ): ...
+    ) -> IntervalArray: ...
     @classmethod
-    def from_tuples(cls, data, closed: str = ..., copy: bool = ..., dtype=...): ...
+    def from_tuples(
+        cls, data, closed: str = ..., copy: bool = ..., dtype=...
+    ) -> IntervalArray: ...
     def __iter__(self): ...
     def __len__(self) -> int: ...
     def __getitem__(self, value): ...

pandas-stubs/core/arrays/period.pyi

@@ -6,10 +6,11 @@ from pandas.core.arrays.datetimelike import (
     DatetimeLikeArrayMixin,
 )
 
-from pandas._libs.tslibs import Timestamp
-from pandas._libs.tslibs.period import Period as Period
+from pandas._libs.tslibs import Period
+from pandas._libs.tslibs.timestamps import Timestamp
+from pandas._typing import npt as npt
 
-from pandas.tseries.offsets import Tick as Tick
+from pandas.tseries.offsets import Tick
 
 class PeriodArray(DatetimeLikeArrayMixin, DatelikeOps):
     __array_priority__: int = ...
@@ -43,11 +44,8 @@ class PeriodArray(DatetimeLikeArrayMixin, DatelikeOps):
     def asfreq(self, freq: str | None = ..., how: str = ...) -> Period: ...
     def astype(self, dtype, copy: bool = ...): ...
 
-def raise_on_incompatible(left, right): ...
 def period_array(
     data: Sequence[Period | None],
     freq: str | Tick | None = ...,
     copy: bool = ...,
 ) -> PeriodArray: ...
-def validate_dtype_freq(dtype, freq): ...
-def dt64arr_to_periodarr(data, freq, tz=...): ...

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

@@ -5,6 +5,8 @@ from pandas.core.arrays import (
 )
 from pandas.core.base import PandasObject
 
+from pandas._libs.sparse import SparseIndex
+
 class SparseArray(PandasObject, ExtensionArray, ExtensionOpsMixin):
     def __init__(
         self,
@@ -20,9 +22,9 @@ class SparseArray(PandasObject, ExtensionArray, ExtensionOpsMixin):
     def __array__(self, dtype=..., copy=...) -> np.ndarray: ...
     def __setitem__(self, key, value) -> None: ...
     @property
-    def sp_index(self): ...
+    def sp_index(self) -> SparseIndex: ...
     @property
-    def sp_values(self): ...
+    def sp_values(self) -> np.ndarray: ...
     @property
     def dtype(self): ...
     @property

pandas-stubs/core/arrays/string_.pyi

@@ -1,6 +1,13 @@
+from typing import Sequence
+
+import numpy as np
+import pandas as pd
 from pandas.core.arrays import PandasArray
 
-from pandas._typing import type_t
+from pandas._typing import (
+    AnyArrayLike,
+    type_t,
+)
 
 from pandas.core.dtypes.base import ExtensionDtype
 
@@ -14,7 +21,12 @@ class StringDtype(ExtensionDtype):
     def __from_arrow__(self, array): ...
 
 class StringArray(PandasArray):
-    def __init__(self, values, copy: bool = ...) -> None: ...
+    def __init__(
+        self,
+        # Also pd.NA and np.nan but not possible it seems
+        values: AnyArrayLike | Sequence[str | None],
+        copy: bool = ...,
+    ) -> None: ...
     def __arrow_array__(self, type=...): ...
     def __setitem__(self, key, value) -> None: ...
     def fillna(self, value=..., method=..., limit=...): ...

pandas-stubs/core/construction.pyi

@@ -1,20 +1,30 @@
-from typing import Sequence
+from typing import (
+    Sequence,
+    Union,
+)
 
 import numpy as np
+import pandas as pd
 from pandas.core.indexes.api import Index
 from pandas.core.series import Series
 
 from pandas._typing import (
     ArrayLike,
     Dtype,
+    npt,
 )
 
 from pandas.core.dtypes.dtypes import ExtensionDtype
 from pandas.core.dtypes.generic import ABCExtensionArray
 
 def array(
-    data: Sequence[object],
-    dtype: str | np.dtype | ExtensionDtype | None = ...,
+    # str is forbidden even though Sequence[object] allows "abc"
+    data: npt.NDArray | Sequence[object] | pd.Index | pd.Series,
+    dtype: str
+    | np.dtype[np.generic]
+    | ExtensionDtype
+    | type[Union[str, bool, float, int]]
+    | None = ...,
     copy: bool = ...,
 ) -> ABCExtensionArray: ...
 def extract_array(obj, extract_numpy: bool = ...): ...