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

@@ -166,6 +166,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")

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[Any]) -> np_1darray[Any]: ...
 @overload
-def unique(values: ExtensionArray) -> ExtensionArray: ...
+def unique(values: T_EXTENSION_ARRAY) -> T_EXTENSION_ARRAY: ...
 @overload
 def factorize(
-    values: np.ndarray,
+    values: np_ndarray[Any, 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: npt.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): ...

pandas-stubs/core/arrays/datetimelike.pyi

@@ -5,7 +5,6 @@ from typing import (
     overload,
 )
 
-import numpy as np
 from pandas.core.arrays.base import (
     ExtensionArray,
     ExtensionOpsMixin,
@@ -65,7 +64,7 @@ class DatetimeLikeArrayMixin(ExtensionOpsMixin, ExtensionArray):
     def ravel(self, *args: Any, **kwargs: Any): ...
     def __iter__(self): ...
     @property
-    def asi8(self) -> np.ndarray: ...
+    def asi8(self) -> np_1darray[Any]: ...
     @property
     def nbytes(self): ...
     def __array__(

pandas-stubs/core/arrays/interval.pyi

@@ -5,6 +5,7 @@ from typing import (
 )
 
 import numpy as np
+from numpy import typing as npt
 from pandas.core.arrays.base import ExtensionArray as ExtensionArray
 from pandas.core.indexes.base import Index
 from pandas.core.series import Series
@@ -105,6 +106,6 @@ class IntervalArray(IntervalMixin, ExtensionArray):
     def contains(self, other: Series) -> Series[bool]: ...
     @overload
     def contains(
-        self, other: Scalar | ExtensionArray | Index | np.ndarray
+        self, other: Scalar | ExtensionArray | Index | npt.NDArray[Any]
     ) -> np_1darray[np.bool]: ...
     def overlaps(self, other: Interval) -> bool: ...

pandas-stubs/core/arrays/masked.pyi

@@ -3,7 +3,6 @@ from typing import (
     overload,
 )
 
-import numpy as np
 from pandas.core.arrays import (
     ExtensionArray as ExtensionArray,
     ExtensionOpsMixin,
@@ -31,7 +30,7 @@ class BaseMaskedArray(ExtensionArray, ExtensionOpsMixin):
         dtype: npt.DTypeLike | None = ...,
         copy: bool = False,
         na_value: Scalar = ...,
-    ) -> np.ndarray: ...
+    ) -> np_1darray[Any]: ...
     __array_priority__: int = ...
     def __array__(
         self, dtype: NpDtype | None = None, copy: bool | None = None