diff options
| -rw-r--r-- | numpy/__init__.pyi | 29 | ||||
| -rw-r--r-- | numpy/core/fromnumeric.pyi | 46 |
2 files changed, 38 insertions, 37 deletions
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 480732c15..fd7731133 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -287,6 +287,21 @@ _OrderKACF = Optional[Literal["K", "A", "C", "F"]] _OrderACF = Optional[Literal["A", "C", "F"]] _OrderCF = Optional[Literal["C", "F"]] +_ModeKind = Literal["raise", "wrap", "clip"] +_PartitionKind = Literal["introselect"] +_SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"] +_SortSide = Literal["left", "right"] + +_ArrayLikeBool = Union[_BoolLike, Sequence[_BoolLike], ndarray] +_ArrayLikeIntOrBool = Union[ + _IntLike, + _BoolLike, + ndarray, + Sequence[_IntLike], + Sequence[_BoolLike], + Sequence[Sequence[Any]], # TODO: wait for support for recursive types +] + _ArraySelf = TypeVar("_ArraySelf", bound=_ArrayOrScalarCommon) class _ArrayOrScalarCommon( @@ -498,11 +513,11 @@ class _ArrayOrScalarCommon( ) -> ndarray: ... @overload def choose( - self, choices: ArrayLike, out: None = ..., mode: _Mode = ..., + self, choices: ArrayLike, out: None = ..., mode: _ModeKind = ..., ) -> ndarray: ... @overload def choose( - self, choices: ArrayLike, out: _NdArraySubClass = ..., mode: _Mode = ..., + self, choices: ArrayLike, out: _NdArraySubClass = ..., mode: _ModeKind = ..., ) -> _NdArraySubClass: ... @overload def clip( @@ -764,7 +779,7 @@ class _ArrayOrScalarCommon( indices: Union[_IntLike, _BoolLike], axis: Optional[int] = ..., out: None = ..., - mode: _Mode = ..., + mode: _ModeKind = ..., ) -> generic: ... @overload def take( @@ -772,7 +787,7 @@ class _ArrayOrScalarCommon( indices: _ArrayLikeIntOrBool, axis: Optional[int] = ..., out: None = ..., - mode: _Mode = ..., + mode: _ModeKind = ..., ) -> ndarray: ... @overload def take( @@ -780,7 +795,7 @@ class _ArrayOrScalarCommon( indices: _ArrayLikeIntOrBool, axis: Optional[int] = ..., out: _NdArraySubClass = ..., - mode: _Mode = ..., + mode: _ModeKind = ..., ) -> _NdArraySubClass: ... @overload def var( @@ -869,12 +884,12 @@ class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container): # `put` is technically available to `generic`, # but is pointless as `generic`s are immutable def put( - self, ind: _ArrayLikeIntOrBool, v: ArrayLike, mode: _Mode = ... + self, ind: _ArrayLikeIntOrBool, v: ArrayLike, mode: _ModeKind = ... ) -> None: ... def searchsorted( self, # >= 1D array v: ArrayLike, - side: _Side = ..., + side: _SortSide = ..., sorter: Optional[_ArrayLikeIntOrBool] = ..., # 1D int array ) -> ndarray: ... def setfield( diff --git a/numpy/core/fromnumeric.pyi b/numpy/core/fromnumeric.pyi index 7ad772b07..6b3d2268f 100644 --- a/numpy/core/fromnumeric.pyi +++ b/numpy/core/fromnumeric.pyi @@ -13,6 +13,12 @@ from numpy import ( _IntLike, _BoolLike, _NumberLike, + _ArrayLikeBool, + _ArrayLikeIntOrBool, + _ModeKind, + _PartitionKind, + _SortKind, + _SortSide, ) from numpy.typing import DtypeLike, ArrayLike, _ShapeLike, _Shape @@ -21,11 +27,6 @@ if sys.version_info >= (3, 8): else: from typing_extensions import Literal -_Mode = Literal["raise", "wrap", "clip"] -_PartitionKind = Literal["introselect"] -_SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"] -_Side = Literal["left", "right"] - # Various annotations for scalars # While dt.datetime and dt.timedelta are not technically part of NumPy, @@ -44,21 +45,6 @@ _ScalarGenericDT = TypeVar( _Number = TypeVar("_Number", bound=number) -# An array-like object consisting of integers -_IntOrBool = Union[_IntLike, _BoolLike] -_ArrayLikeIntNested = ArrayLike # TODO: wait for support for recursive types -_ArrayLikeBoolNested = ArrayLike # TODO: wait for support for recursive types - -# Integers and booleans can generally be used interchangeably -_ArrayLikeIntOrBool = Union[ - _IntOrBool, - ndarray, - Sequence[_IntOrBool], - Sequence[_ArrayLikeIntNested], - Sequence[_ArrayLikeBoolNested], -] -_ArrayLikeBool = Union[_BoolLike, Sequence[_BoolLike], ndarray] - # The signature of take() follows a common theme with its overloads: # 1. A generic comes in; the same generic comes out # 2. A scalar comes in; a generic comes out @@ -70,7 +56,7 @@ def take( indices: int, axis: Optional[int] = ..., out: Optional[ndarray] = ..., - mode: _Mode = ..., + mode: _ModeKind = ..., ) -> _ScalarGenericDT: ... @overload def take( @@ -78,7 +64,7 @@ def take( indices: int, axis: Optional[int] = ..., out: Optional[ndarray] = ..., - mode: _Mode = ..., + mode: _ModeKind = ..., ) -> _ScalarNumpy: ... @overload def take( @@ -86,7 +72,7 @@ def take( indices: int, axis: Optional[int] = ..., out: Optional[ndarray] = ..., - mode: _Mode = ..., + mode: _ModeKind = ..., ) -> _ScalarNumpy: ... @overload def take( @@ -94,7 +80,7 @@ def take( indices: _ArrayLikeIntOrBool, axis: Optional[int] = ..., out: Optional[ndarray] = ..., - mode: _Mode = ..., + mode: _ModeKind = ..., ) -> Union[_ScalarNumpy, ndarray]: ... def reshape(a: ArrayLike, newshape: _ShapeLike, order: _OrderACF = ...) -> ndarray: ... @overload @@ -102,24 +88,24 @@ def choose( a: _ScalarIntOrBool, choices: ArrayLike, out: Optional[ndarray] = ..., - mode: _Mode = ..., + mode: _ModeKind = ..., ) -> _ScalarIntOrBool: ... @overload def choose( - a: _IntOrBool, choices: ArrayLike, out: Optional[ndarray] = ..., mode: _Mode = ... + a: Union[_IntLike, _BoolLike], choices: ArrayLike, out: Optional[ndarray] = ..., mode: _ModeKind = ... ) -> Union[integer, bool_]: ... @overload def choose( a: _ArrayLikeIntOrBool, choices: ArrayLike, out: Optional[ndarray] = ..., - mode: _Mode = ..., + mode: _ModeKind = ..., ) -> ndarray: ... def repeat( a: ArrayLike, repeats: _ArrayLikeIntOrBool, axis: Optional[int] = ... ) -> ndarray: ... def put( - a: ndarray, ind: _ArrayLikeIntOrBool, v: ArrayLike, mode: _Mode = ... + a: ndarray, ind: _ArrayLikeIntOrBool, v: ArrayLike, mode: _ModeKind = ... ) -> None: ... def swapaxes(a: ArrayLike, axis1: int, axis2: int) -> ndarray: ... def transpose( @@ -184,14 +170,14 @@ def argmin( def searchsorted( a: ArrayLike, v: _Scalar, - side: _Side = ..., + side: _SortSide = ..., sorter: Optional[_ArrayLikeIntOrBool] = ..., # 1D int array ) -> integer: ... @overload def searchsorted( a: ArrayLike, v: ArrayLike, - side: _Side = ..., + side: _SortSide = ..., sorter: Optional[_ArrayLikeIntOrBool] = ..., # 1D int array ) -> ndarray: ... def resize(a: ArrayLike, new_shape: _ShapeLike) -> ndarray: ... |