diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index c999e17..1d19b6e 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -772,6 +772,7 @@ _Mode = Literal["raise", "wrap", "clip"] _Order = Literal["C", "F", "A"] _PartitionKind = Literal["introselect"] _SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"] +_Side = Literal["left", "right"] # Various annotations for scalars @@ -793,8 +794,8 @@ _ScalarGenericDT = TypeVar( _Int = Union[int, integer] _Bool = Union[bool, bool_] _IntOrBool = Union[_Int, _Bool] -_ArrayLikeIntNested = Any # TODO: wait for support for recursive types -_ArrayLikeBoolNested = Any # TODO: wait for support for recursive types +_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[ @@ -919,3 +920,69 @@ def argsort( kind: Optional[_SortKind] = ..., order: Union[None, str, Sequence[str]] = ..., ) -> ndarray: ... +@overload +def argmax( + a: Union[Sequence[ArrayLike], ndarray], + axis: None = ..., + out: Optional[ndarray] = ..., +) -> integer: ... +@overload +def argmax( + a: Union[Sequence[ArrayLike], ndarray], + axis: int = ..., + out: Optional[ndarray] = ..., +) -> Union[integer, ndarray]: ... +@overload +def argmin( + a: Union[Sequence[ArrayLike], ndarray], + axis: None = ..., + out: Optional[ndarray] = ..., +) -> integer: ... +@overload +def argmin( + a: Union[Sequence[ArrayLike], ndarray], + axis: int = ..., + out: Optional[ndarray] = ..., +) -> Union[integer, ndarray]: ... +@overload +def searchsorted( + a: Union[Sequence[ArrayLike], ndarray], + v: _Scalar, + side: _Side = ..., + sorter: Union[None, Sequence[_IntOrBool], ndarray] = ..., # 1D int array +) -> integer: ... +@overload +def searchsorted( + a: Union[Sequence[ArrayLike], ndarray], + v: ArrayLike, + side: _Side = ..., + sorter: Union[None, Sequence[_IntOrBool], ndarray] = ..., # 1D int array +) -> ndarray: ... +def resize(a: ArrayLike, new_shape: _ShapeLike) -> ndarray: ... +@overload +def squeeze(a: _ScalarGeneric, axis: Optional[_ShapeLike] = ...) -> _ScalarGeneric: ... +@overload +def squeeze(a: ArrayLike, axis: Optional[_ShapeLike] = ...) -> ndarray: ... +def diagonal( + a: Union[Sequence[Sequence[ArrayLike]], ndarray], # >= 2D array + offset: int = ..., + axis1: int = ..., + axis2: int = ..., +) -> ndarray: ... +def trace( + a: Union[Sequence[Sequence[ArrayLike]], ndarray], # >= 2D array + offset: int = ..., + axis1: int = ..., + axis2: int = ..., + dtype: DtypeLike = ..., + out: Optional[ndarray] = ..., +) -> Union[number, ndarray]: ... +def ravel(a: ArrayLike, order: _Order = ...) -> ndarray: ... +def nonzero(a: ArrayLike) -> Tuple[ndarray, ...]: ... +def shape(a: ArrayLike) -> _Shape: ... +def compress( + condition: Union[Sequence[_Bool], ndarray], # 1D bool array + a: ArrayLike, + axis: Optional[int] = ..., + out: Optional[ndarray] = ..., +) -> ndarray: ... diff --git a/tests/fail/fromnumeric.py b/tests/fail/fromnumeric.py index f5f6fdb..f158a10 100644 --- a/tests/fail/fromnumeric.py +++ b/tests/fail/fromnumeric.py @@ -62,3 +62,40 @@ np.argsort(A, axis="bob") # E: Argument "axis" to "argsort" has incompatible type np.argsort(A, kind="bob") # E: Argument "kind" to "argsort" has incompatible type np.argsort(A, order=range(5)) # E: Argument "order" to "argsort" has incompatible type + +np.argmax(a) # E: No overload variant of "argmax" matches argument type +np.argmax(A, axis="bob") # E: No overload variant of "argmax" matches argument type +np.argmax(A, kind="bob") # E: No overload variant of "argmax" matches argument type + +np.argmin(a) # E: No overload variant of "argmin" matches argument type +np.argmin(A, axis="bob") # E: No overload variant of "argmin" matches argument type +np.argmin(A, kind="bob") # E: No overload variant of "argmin" matches argument type + +np.searchsorted(a, 0) # E: No overload variant of "searchsorted" matches argument type +np.searchsorted( # E: No overload variant of "searchsorted" matches argument type + A[0], 0, side="bob" +) +np.searchsorted( # E: No overload variant of "searchsorted" matches argument type + A[0], 0, sorter=1.0 +) + +np.resize(A, 1.0) # E: Argument 2 to "resize" has incompatible type + +np.squeeze(A, 1.0) # E: No overload variant of "squeeze" matches argument type + +np.diagonal(a) # E: Argument 1 to "diagonal" has incompatible type +np.diagonal(A, offset=None) # E: Argument "offset" to "diagonal" has incompatible type +np.diagonal(A, axis1="bob") # E: Argument "axis1" to "diagonal" has incompatible type +np.diagonal(A, axis2=[]) # E: Argument "axis2" to "diagonal" has incompatible type + +np.trace(a) # E: Argument 1 to "trace" has incompatible type +np.trace(A, offset=None) # E: Argument "offset" to "trace" has incompatible type +np.trace(A, axis1="bob") # E: Argument "axis1" to "trace" has incompatible type +np.trace(A, axis2=[]) # E: Argument "axis2" to "trace" has incompatible type + +np.ravel(a, order="bob") # E: Argument "order" to "ravel" has incompatible type + +np.compress(True, A) # E: Argument 1 to "compress" has incompatible type +np.compress( + [True], A, axis=1.0 # E: Argument "axis" to "compress" has incompatible type +) diff --git a/tests/pass/fromnumeric.py b/tests/pass/fromnumeric.py index 4a97049..a66864b 100644 --- a/tests/pass/fromnumeric.py +++ b/tests/pass/fromnumeric.py @@ -60,3 +60,60 @@ np.argsort(A, 0) np.argsort(B, 0) + +np.argmax(A) +np.argmax(B) +np.argmax(A, axis=0) +np.argmax(B, axis=0) + +np.argmin(A) +np.argmin(B) +np.argmin(A, axis=0) +np.argmin(B, axis=0) + +np.searchsorted(A[0], 0) +np.searchsorted(B[0], 0) +np.searchsorted(A[0], [0]) +np.searchsorted(B[0], [0]) + +np.resize(a, (5, 5)) +np.resize(b, (5, 5)) +np.resize(c, (5, 5)) +np.resize(A, (5, 5)) +np.resize(B, (5, 5)) + +np.squeeze(a) +np.squeeze(b) +np.squeeze(c) +np.squeeze(A) +np.squeeze(B) + +np.diagonal(A) +np.diagonal(B) + +np.trace(A) +np.trace(B) + +np.ravel(a) +np.ravel(b) +np.ravel(c) +np.ravel(A) +np.ravel(B) + +np.nonzero(a) +np.nonzero(b) +np.nonzero(c) +np.nonzero(A) +np.nonzero(B) + +np.shape(a) +np.shape(b) +np.shape(c) +np.shape(A) +np.shape(B) + +np.compress([True], a) +np.compress([True], b) +np.compress([True], c) +np.compress([True], A) +np.compress([True], B) diff --git a/tests/reveal/fromnumeric.py b/tests/reveal/fromnumeric.py index 1fde52c..7d79d5f 100644 --- a/tests/reveal/fromnumeric.py +++ b/tests/reveal/fromnumeric.py @@ -76,3 +76,60 @@ reveal_type(np.argsort(A, 0)) # E: numpy.ndarray reveal_type(np.argsort(B, 0)) # E: numpy.ndarray + +reveal_type(np.argmax(A)) # E: numpy.integer +reveal_type(np.argmax(B)) # E: numpy.integer +reveal_type(np.argmax(A, axis=0)) # E: Union[numpy.integer, numpy.ndarray] +reveal_type(np.argmax(B, axis=0)) # E: Union[numpy.integer, numpy.ndarray] + +reveal_type(np.argmin(A)) # E: numpy.integer +reveal_type(np.argmin(B)) # E: numpy.integer +reveal_type(np.argmin(A, axis=0)) # E: Union[numpy.integer, numpy.ndarray] +reveal_type(np.argmin(B, axis=0)) # E: Union[numpy.integer, numpy.ndarray] + +reveal_type(np.searchsorted(A[0], 0)) # E: numpy.integer +reveal_type(np.searchsorted(B[0], 0)) # E: numpy.integer +reveal_type(np.searchsorted(A[0], [0])) # E: numpy.ndarray +reveal_type(np.searchsorted(B[0], [0])) # E: numpy.ndarray + +reveal_type(np.resize(a, (5, 5))) # E: numpy.ndarray +reveal_type(np.resize(b, (5, 5))) # E: numpy.ndarray +reveal_type(np.resize(c, (5, 5))) # E: numpy.ndarray +reveal_type(np.resize(A, (5, 5))) # E: numpy.ndarray +reveal_type(np.resize(B, (5, 5))) # E: numpy.ndarray + +reveal_type(np.squeeze(a)) # E: numpy.bool_ +reveal_type(np.squeeze(b)) # E: numpy.float32 +reveal_type(np.squeeze(c)) # E: numpy.ndarray +reveal_type(np.squeeze(A)) # E: numpy.ndarray +reveal_type(np.squeeze(B)) # E: numpy.ndarray + +reveal_type(np.diagonal(A)) # E: numpy.ndarray +reveal_type(np.diagonal(B)) # E: numpy.ndarray + +reveal_type(np.trace(A)) # E: Union[numpy.number, numpy.ndarray] +reveal_type(np.trace(B)) # E: Union[numpy.number, numpy.ndarray] + +reveal_type(np.ravel(a)) # E: numpy.ndarray +reveal_type(np.ravel(b)) # E: numpy.ndarray +reveal_type(np.ravel(c)) # E: numpy.ndarray +reveal_type(np.ravel(A)) # E: numpy.ndarray +reveal_type(np.ravel(B)) # E: numpy.ndarray + +reveal_type(np.nonzero(a)) # E: tuple[numpy.ndarray] +reveal_type(np.nonzero(b)) # E: tuple[numpy.ndarray] +reveal_type(np.nonzero(c)) # E: tuple[numpy.ndarray] +reveal_type(np.nonzero(A)) # E: tuple[numpy.ndarray] +reveal_type(np.nonzero(B)) # E: tuple[numpy.ndarray] + +reveal_type(np.shape(a)) # E: tuple[builtins.int] +reveal_type(np.shape(b)) # E: tuple[builtins.int] +reveal_type(np.shape(c)) # E: tuple[builtins.int] +reveal_type(np.shape(A)) # E: tuple[builtins.int] +reveal_type(np.shape(B)) # E: tuple[builtins.int] + +reveal_type(np.compress([True], a)) # E: numpy.ndarray +reveal_type(np.compress([True], b)) # E: numpy.ndarray +reveal_type(np.compress([True], c)) # E: numpy.ndarray +reveal_type(np.compress([True], A)) # E: numpy.ndarray +reveal_type(np.compress([True], B)) # E: numpy.ndarray