reduce ndarray

Author: cmp0xff

pandas-stubs/_libs/interval.pyi

@@ -4,6 +4,7 @@ from typing import (
     Literal,
     TypeVar,
     overload,
+    type_check_only,
 )
 
 import numpy as np
@@ -18,7 +19,6 @@ from pandas._typing import (
     IntervalClosedType,
     IntervalT,
     np_1darray,
-    npt,
 )
 
 VALID_CLOSED: frozenset[str]
@@ -27,6 +27,7 @@ _OrderableScalarT = TypeVar("_OrderableScalarT", bound=int | float)
 _OrderableTimesT = TypeVar("_OrderableTimesT", bound=Timestamp | Timedelta)
 _OrderableT = TypeVar("_OrderableT", bound=int | float | Timestamp | Timedelta)
 
+@type_check_only
 class _LengthDescriptor:
     @overload
     def __get__(
@@ -36,18 +37,15 @@ class _LengthDescriptor:
     def __get__(
         self, instance: Interval[_OrderableTimesT], owner: Any
     ) -> Timedelta: ...
-    @overload
-    def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ...
 
+@type_check_only
 class _MidDescriptor:
     @overload
     def __get__(self, instance: Interval[_OrderableScalarT], owner: Any) -> float: ...
     @overload
     def __get__(
         self, instance: Interval[_OrderableTimesT], owner: Any
     ) -> _OrderableTimesT: ...
-    @overload
-    def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ...
 
 class IntervalMixin:
     @property
@@ -68,8 +66,8 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     def right(self: Interval[_OrderableT]) -> _OrderableT: ...
     @property
     def closed(self) -> IntervalClosedType: ...
-    mid: _MidDescriptor
-    length: _LengthDescriptor
+    mid = _MidDescriptor()
+    length = _LengthDescriptor()
     def __init__(
         self,
         left: _OrderableT,
@@ -223,21 +221,4 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     @overload
     def __ne__(self, other: object) -> Literal[True]: ...
 
-class IntervalTree(IntervalMixin):
-    def __init__(
-        self,
-        left: np.ndarray,
-        right: np.ndarray,
-        closed: IntervalClosedType = ...,
-        leaf_size: int = ...,
-    ) -> None: ...
-    def get_indexer(self, target) -> npt.NDArray[np.intp]: ...
-    def get_indexer_non_unique(
-        self, target
-    ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ...
-    _na_count: int
-    @property
-    def is_overlapping(self) -> bool: ...
-    @property
-    def is_monotonic_increasing(self) -> bool: ...
-    def clear_mapping(self) -> None: ...
+class IntervalTree(IntervalMixin): ...

pandas-stubs/_libs/tslibs/offsets.pyi

@@ -14,10 +14,14 @@ from typing import (
 
 from dateutil.relativedelta import weekday as WeekdayClass
 import numpy as np
+from numpy import typing as npt
 from pandas import Timestamp
 from typing_extensions import Self
 
-from pandas._typing import npt
+from pandas._typing import (
+    ShapeT,
+    np_ndarray,
+)
 
 from pandas.tseries.holiday import AbstractHolidayCalendar
 
@@ -37,7 +41,9 @@ class BaseOffset:
     @property
     def base(self) -> BaseOffset: ...
     @overload
-    def __add__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
+    def __add__(
+        self, other: np_ndarray[ShapeT, np.object_]
+    ) -> np_ndarray[ShapeT, np.object_]: ...
     @overload
     def __add__(self, other: _DatetimeT) -> _DatetimeT: ...  # type: ignore[overload-overlap]  # pyright: ignore[reportOverlappingOverload]
     @overload
@@ -47,7 +53,9 @@ class BaseOffset:
     @overload
     def __add__(self, other: _TimedeltaT) -> _TimedeltaT: ...
     @overload
-    def __radd__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
+    def __radd__(
+        self, other: np_ndarray[ShapeT, np.object_]
+    ) -> np_ndarray[ShapeT, np.object_]: ...
     @overload
     def __radd__(self, other: _DatetimeT) -> _DatetimeT: ...  # type: ignore[overload-overlap]  # pyright: ignore[reportOverlappingOverload]
     @overload
@@ -68,11 +76,15 @@ class BaseOffset:
     @overload
     def __rsub__(self, other: _TimedeltaT) -> _TimedeltaT: ...
     @overload
-    def __mul__(self, other: np.ndarray) -> np.ndarray: ...
+    def __mul__(
+        self, other: np_ndarray[ShapeT, np.object_]
+    ) -> np_ndarray[ShapeT, np.object_]: ...
     @overload
     def __mul__(self, other: int) -> Self: ...
     @overload
-    def __rmul__(self, other: np.ndarray) -> np.ndarray: ...
+    def __rmul__(
+        self, other: np_ndarray[ShapeT, np.object_]
+    ) -> np_ndarray[ShapeT, np.object_]: ...
     @overload
     def __rmul__(self, other: int) -> Self: ...
     def __neg__(self) -> Self: ...

pandas-stubs/_testing/__init__.pyi

@@ -12,7 +12,7 @@ from typing import (
 import warnings
 
 from matplotlib.artist import Artist
-import numpy as np
+from numpy import typing as npt
 from pandas import (
     Categorical,
     DataFrame,
@@ -61,7 +61,7 @@ def assert_attr_equal(
     attr: str, left: object, right: object, obj: str = "Attributes"
 ) -> None: ...
 def assert_is_valid_plot_return_object(
-    objs: Series | np.ndarray | Artist | tuple | dict,
+    objs: Series | npt.NDArray[Any] | Artist | tuple | dict,
 ) -> None: ...
 def assert_is_sorted(seq: AnyArrayLike) -> None: ...
 def assert_categorical_equal(
@@ -96,7 +96,7 @@ def assert_extension_array_equal(
     left: ExtensionArray,
     right: ExtensionArray,
     check_dtype: bool | Literal["equiv"] = True,
-    index_values: Index | np.ndarray | None = None,
+    index_values: Index | npt.NDArray[Any] | None = None,
     check_exact: bool = False,
     rtol: float = 1e-5,
     atol: float = 1e-8,

pandas-stubs/_typing.pyi

@@ -85,10 +85,13 @@ HashableT5 = TypeVar("HashableT5", bound=Hashable)
 
 # array-like
 
-ArrayLike: TypeAlias = ExtensionArray | np.ndarray
+ArrayLike: TypeAlias = ExtensionArray | np.ndarray[tuple[int], np.dtype[Any]]
 AnyArrayLike: TypeAlias = ArrayLike | Index | Series
 AnyArrayLikeInt: TypeAlias = (
-    IntegerArray | Index[int] | Series[int] | npt.NDArray[np.integer]
+    IntegerArray
+    | Index[int]
+    | Series[int]
+    | np.ndarray[tuple[int], np.dtype[np.integer]]
 )
 
 # list-like
@@ -166,6 +169,7 @@ ToTimestampHow: TypeAlias = Literal["s", "e", "start", "end"]
 NDFrameT = TypeVar("NDFrameT", bound=NDFrame)
 
 IndexT = TypeVar("IndexT", bound=Index)
+T_EXTENSION_ARRAY = TypeVar("T_EXTENSION_ARRAY", bound=ExtensionArray)
 
 # From _typing.py, not used here:
 # FreqIndexT = TypeVar("FreqIndexT", "DatetimeIndex", "PeriodIndex", "TimedeltaIndex")
@@ -697,11 +701,11 @@ InterpolateOptions: TypeAlias = Literal[
 # Using List[int] here rather than Sequence[int] to disallow tuples.
 
 ScalarIndexer: TypeAlias = int | np.integer
-SequenceIndexer: TypeAlias = slice | list[int] | np.ndarray
+SequenceIndexer: TypeAlias = slice | list[int] | npt.NDArray[np.integer | np.bool]
 PositionalIndexer: TypeAlias = ScalarIndexer | SequenceIndexer
 PositionalIndexerTuple: TypeAlias = tuple[PositionalIndexer, PositionalIndexer]
 # PositionalIndexer2D = Union[PositionalIndexer, PositionalIndexerTuple] Not used in stubs
-TakeIndexer: TypeAlias = Sequence[int] | Sequence[np.integer] | npt.NDArray[np.integer]
+TakeIndexer: TypeAlias = Sequence[int | np.integer] | npt.NDArray[np.integer | np.bool]
 
 # Shared by functions such as drop and astype
 IgnoreRaise: TypeAlias = Literal["ignore", "raise"]
@@ -832,18 +836,6 @@ SliceType: TypeAlias = Hashable | None
 ## All types below this point are only used in pandas-stubs
 ######
 
-num: TypeAlias = complex
-
-DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic])
-KeysArgType: TypeAlias = Any
-ListLikeT = TypeVar("ListLikeT", bound=ListLike)
-ListLikeExceptSeriesAndStr: TypeAlias = (
-    MutableSequence[Any] | np.ndarray | tuple[Any, ...] | Index
-)
-ListLikeU: TypeAlias = Sequence | np.ndarray | Series | Index
-ListLikeHashable: TypeAlias = (
-    MutableSequence[HashableT] | np.ndarray | tuple[HashableT, ...] | range
-)
 StrLike: TypeAlias = str | np.str_
 
 ScalarT = TypeVar("ScalarT", bound=Scalar)
@@ -871,6 +863,17 @@ np_ndarray: TypeAlias = np.ndarray[ShapeT, np.dtype[GenericT]]
 np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[GenericT]]
 np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]]
 
+DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic])
+KeysArgType: TypeAlias = Any
+ListLikeT = TypeVar("ListLikeT", bound=ListLike)
+ListLikeExceptSeriesAndStr: TypeAlias = (
+    MutableSequence[Any] | np_1darray[Any] | tuple[Any, ...] | Index
+)
+ListLikeU: TypeAlias = Sequence | np_1darray[Any] | Series | Index
+ListLikeHashable: TypeAlias = (
+    MutableSequence[HashableT] | np_1darray[Any] | tuple[HashableT, ...] | range
+)
+
 class SupportsDType(Protocol[GenericT_co]):
     @property
     def dtype(self) -> np.dtype[GenericT_co]: ...

pandas-stubs/core/algorithms.pyi

@@ -5,65 +5,75 @@ from typing import (
 )
 
 import numpy as np
-from pandas import (
-    Categorical,
-    CategoricalIndex,
-    Index,
-    IntervalIndex,
-    PeriodIndex,
-    Series,
-)
+from numpy import typing as npt
 from pandas.api.extensions import ExtensionArray
+from pandas.core.arrays.categorical import Categorical
+from pandas.core.indexes.base import Index
+from pandas.core.indexes.category import CategoricalIndex
+from pandas.core.indexes.datetimes import DatetimeIndex
+from pandas.core.indexes.interval import IntervalIndex
+from pandas.core.indexes.period import PeriodIndex
+from pandas.core.series import Series
 
 from pandas._typing import (
+    T_EXTENSION_ARRAY,
     AnyArrayLike,
+    GenericT,
     IntervalT,
     TakeIndexer,
     np_1darray,
+    np_ndarray,
 )
 
 # These are type: ignored because the Index types overlap due to inheritance but indices
 # with extension types return the same type while standard type return ndarray
-
-@overload
-def unique(  # pyright: ignore[reportOverlappingOverload]
-    values: PeriodIndex,
-) -> PeriodIndex: ...
 @overload
-def unique(values: CategoricalIndex) -> CategoricalIndex: ...  # type: ignore[overload-overlap]
+def unique(values: CategoricalIndex) -> CategoricalIndex: ...
 @overload
 def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ...
 @overload
-def unique(values: Index) -> np.ndarray: ...
+def unique(values: PeriodIndex) -> PeriodIndex: ...
+@overload
+def unique(values: DatetimeIndex) -> np_1darray[np.datetime64] | DatetimeIndex: ...
+@overload
+def unique(values: Index) -> np_1darray[Any] | Index: ...
 @overload
 def unique(values: Categorical) -> Categorical: ...
+
+# @overload
+# def unique(values: Series[Never]) -> np_1darray[Any] | ExtensionArray: ...
+# TODO: DatetimeArray python/mypy#19952
+# @overload
+# def unique(values: Series[Timestamp]) -> np_1darray[np.datetime64] | ExtensionArray: ...
+# @overload
+# def unique(values: Series[int]) -> np_1darray[np.integer] | ExtensionArray: ...
 @overload
-def unique(values: Series) -> np.ndarray | ExtensionArray: ...
+def unique(values: Series) -> np_1darray[Any] | ExtensionArray: ...
 @overload
-def unique(values: np.ndarray) -> np.ndarray: ...
+def unique(values: npt.NDArray[GenericT]) -> np_1darray[GenericT]: ...
 @overload
-def unique(values: ExtensionArray) -> ExtensionArray: ...
+def unique(values: T_EXTENSION_ARRAY) -> T_EXTENSION_ARRAY: ...
 @overload
 def factorize(
-    values: np.ndarray,
+    values: npt.NDArray[GenericT],
     sort: bool = ...,
     use_na_sentinel: bool = ...,
     size_hint: int | None = ...,
-) -> tuple[np.ndarray, np.ndarray]: ...
+) -> tuple[np_1darray[np.int64], np_1darray[GenericT]]: ...
 @overload
 def factorize(
     values: Index | Series,
     sort: bool = ...,
     use_na_sentinel: bool = ...,
     size_hint: int | None = ...,
-) -> tuple[np_1darray, Index]: ...
+) -> tuple[np_1darray[np.int64], Index]: ...
 @overload
 def factorize(
     values: Categorical,
     sort: bool = ...,
     use_na_sentinel: bool = ...,
     size_hint: int | None = ...,
-) -> tuple[np_1darray, Categorical]: ...
+) -> tuple[np_1darray[np.int64], Categorical]: ...
 def value_counts(
     values: AnyArrayLike | list | tuple,
     sort: bool = True,
@@ -73,9 +83,9 @@ def value_counts(
     dropna: bool = True,
 ) -> Series: ...
 def take(
-    arr: np.ndarray | ExtensionArray | Index | Series,
+    arr: np_ndarray[Any] | ExtensionArray | Index | Series,
     indices: TakeIndexer,
     axis: Literal[0, 1] = 0,
     allow_fill: bool = False,
     fill_value: Any = None,
-) -> np_1darray | ExtensionArray: ...
+) -> np_1darray[Any] | ExtensionArray: ...

pandas-stubs/core/arrays/base.pyi

@@ -35,7 +35,7 @@ class ExtensionArray:
     def __getitem__(self, item: ScalarIndexer) -> Any: ...
     @overload
     def __getitem__(self, item: SequenceIndexer) -> Self: ...
-    def __setitem__(self, key: int | slice | np.ndarray, value: Any) -> None: ...
+    def __setitem__(self, key: int | slice | npt.NDArray[Any], value: Any) -> None: ...
     def __len__(self) -> int: ...
     def __iter__(self) -> Iterator[Any]: ...
     def __contains__(self, item: object) -> bool | np.bool_: ...
@@ -83,7 +83,9 @@ class ExtensionArray:
         side: Literal["left", "right"] = ...,
         sorter: ListLike | None = ...,
     ) -> np.intp: ...
-    def factorize(self, use_na_sentinel: bool = True) -> tuple[np_1darray, Self]: ...
+    def factorize(
+        self, use_na_sentinel: bool = True
+    ) -> tuple[np_1darray[Any], Self]: ...
     def repeat(
         self, repeats: int | AnyArrayLikeInt | Sequence[int], axis: None = None
     ) -> Self: ...

pandas-stubs/core/arrays/boolean.pyi

@@ -1,10 +1,12 @@
 from typing import Any
 
-import numpy as np
 from pandas.core.arrays.masked import BaseMaskedArray as BaseMaskedArray
 
 from pandas._libs.missing import NAType
-from pandas._typing import type_t
+from pandas._typing import (
+    np_ndarray_bool,
+    type_t,
+)
 
 from pandas.core.dtypes.base import ExtensionDtype as ExtensionDtype
 
@@ -16,7 +18,7 @@ class BooleanDtype(ExtensionDtype):
 
 class BooleanArray(BaseMaskedArray):
     def __init__(
-        self, values: np.ndarray, mask: np.ndarray, copy: bool = ...
+        self, values: np_ndarray_bool, mask: np_ndarray_bool, copy: bool = ...
     ) -> None: ...
     @property
     def dtype(self): ...