ENH: More array improvements

Author: bashtage

pandas-stubs/core/arrays/boolean.pyi

@@ -48,7 +48,10 @@ class BooleanArray(BaseMaskedArray):
     def __setitem__(
         self,
         key: _ArrayKey,
-        value: _ScalarType | Sequence[bool | NAType | None] | npt.NDArray[np.bool_],
+        value: _ScalarType
+        | Sequence[bool | NAType | None]
+        | npt.NDArray[np.bool_]
+        | BooleanArray,
     ) -> None: ...
     @overload  # type: ignore[override]
     def __getitem__(self, item: int) -> bool | NAType: ...

pandas-stubs/core/arrays/string_.pyi

@@ -1,7 +1,18 @@
+from typing import (
+    Sequence,
+    Union,
+    overload,
+)
+
 import numpy as np
 import pandas as pd
-from pandas.core.arrays import PandasArray
+from pandas.core.arrays import (
+    ExtensionArray,
+    PandasArray,
+)
+from typing_extensions import TypeAlias
 
+from pandas._libs.missing import NAType
 from pandas._typing import (
     npt,
     type_t,
@@ -18,6 +29,9 @@ class StringDtype(ExtensionDtype):
     def construct_array_type(cls) -> type_t[StringArray]: ...
     def __from_arrow__(self, array): ...
 
+_ScalarType: TypeAlias = Union[str, NAType, None]
+_ArrayKey: TypeAlias = Union[Sequence[int], npt.NDArray[np.integer], slice]
+
 class StringArray(PandasArray):
     def __init__(
         self,
@@ -26,8 +40,34 @@ class StringArray(PandasArray):
         | npt.NDArray[np.object_],
         copy: bool = ...,
     ) -> None: ...
+    @property
+    def na_value(self) -> NAType: ...
     def __arrow_array__(self, type=...): ...
-    def __setitem__(self, key, value) -> None: ...
-    def fillna(self, value=..., method=..., limit=...): ...
-    def astype(self, dtype, copy: bool = ...): ...
-    def value_counts(self, dropna: bool = ...): ...
+    @overload  # type: ignore[override]
+    def __setitem__(self, key: int, value: _ScalarType) -> None: ...
+    @overload
+    def __setitem__(
+        self,
+        key: _ArrayKey,
+        value: _ScalarType
+        | Sequence[str | NAType | None]
+        | npt.NDArray[np.str_]
+        | npt.NDArray[np.string_]
+        | StringArray,
+    ) -> None: ...
+    @overload  # type: ignore[override]
+    def __getitem__(self, item: int) -> str | NAType: ...
+    @overload
+    def __getitem__(self, item: _ArrayKey) -> StringArray: ...
+    @overload
+    def astype(self, dtype: npt.DTypeLike, copy: bool = ...) -> np.ndarray: ...
+    @overload
+    def astype(self, dtype: ExtensionDtype, copy: bool = ...) -> ExtensionArray: ...
+    def min(
+        self, axis: int | None = ..., skipna: bool = ..., **kwargs
+    ) -> str | NAType: ...
+    def max(
+        self, axis: int | None = ..., skipna: bool = ..., **kwargs
+    ) -> str | NAType: ...
+    def value_counts(self, dropna: bool = ...) -> pd.Series[int]: ...
+    def memory_usage(self, deep: bool = ...) -> int: ...

pandas-stubs/core/arrays/timedeltas.pyi

@@ -63,7 +63,3 @@ class TimedeltaArray(DatetimeLikeArrayMixin, TimelikeOps):
     nanoseconds: int = ...
     @property
     def components(self) -> int: ...
-
-def sequence_to_td64ns(data, copy: bool = ..., unit: str = ..., errors: str = ...): ...
-def ints_to_td64ns(data, unit: str = ...): ...
-def objects_to_td64ns(data, unit: str = ..., errors: str = ...): ...

tests/test_arrays.py

@@ -499,6 +499,64 @@ def test_timedelta_array() -> None:
     check(assert_type(pd.array(tdi), Type[ExtensionArray]), TimedeltaArray)
     check(assert_type(pd.array(tds, copy=False), Type[ExtensionArray]), TimedeltaArray)
 
+    # check(assert_type(tda.microseconds,npt.NDArray[np.int_]), np.ndarray)
+    # check(assert_type(tda.nanoseconds,npt.NDArray[np.int_]), np.ndarray)
+    # check(assert_type(tda.days,npt.NDArray[np.int_]), np.ndarray)
+    # check(assert_type(tda.seconds,npt.NDArray[np.int_]), np.ndarray)
+    #
+    # all
+    # any
+    # check(assert_type(tda.argmax(), int), int)
+    # check(assert_type(tda.argmin(), int), int)
+    # check(assert_type(argsort, npt.NDArray[npt.intp]), np.ndarray)
+    # check(assert_type(tda.asi8, npt.NDArray[npt.intp]), np.ndarray)
+    #
+    # astype
+    # ceil
+    # components
+    # copy
+    # delete
+    # dropna
+    # dtype
+    # equals
+    # factorize
+    # fillna
+    # floor
+    # freq
+    # freqstr
+    # inferred_freq
+    # insert
+    # isin
+    # isna
+    # map
+    # max
+    # mean
+    # median
+    # min
+    # nbytes
+    # ndim
+    # ravel
+    # repeat
+    # reshape
+    # resolution
+    # round
+    # searchsorted
+    # shape
+    # shift
+    # size
+    # std
+    # sum
+    # swapaxes
+    # take
+    # check(assert_type(tda.to_numpy(), npt.NDArray[np.timedelta64]), np.ndarray)
+    # check(assert_type(tda.to_pytimedelta(), npt.NDArray[np.object_]), np.ndarray)
+    # tsa.tolist()
+    # check(assert_type(tda.total_seconds(), npt.NDArray[np.float64]),np.ndarray)
+    # transpose
+    # unique
+    # value_counts
+    # view
+
 
 def test_sparse_array() -> None:
     ints = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]