From 689dc2dfefb6fb500b46f567993b6492bae0d5f9 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Fri, 24 Apr 2020 15:05:31 +0200 Subject: [PATCH 1/3] Added 11 new functions + tests and updated some existing annotations New functions: * ``argmax()`` * ``argmin()`` * ``searchsorted()`` * ``resize()`` * ``squeeze()`` * ``diagonal()`` * ``trace()`` * ``ravel()`` * ``nonzero()`` * ``shape()`` * ``compress()`` Updates: * Integers and booleans can often be used interchangebly (in fact, ``bool`` is an ``int`` subclass in builtin python). Updated ``_ArrayLikeInt`` to reflect this compatibility. * Rename ``_ScalarGeneric`` to ``_ScalarGenericPlus``; add a new ``_ScalarGeneric`` TypeVar which is *only* bound to ``np.generic``. * ``np.choose()`` cannot only accept integers (or booleans); reflect this in its annotation. * It seems ``np.partition()`` is bugged in NumPy 1.17.3; it's default value for ``axis`` is set to ``-1`` instead of ``None``. Consquently, it will raise an ``AxisError`` if a scalar is passed *unless* one explicitly passes ``axis=None`` to the function. * ``np.argpartition()`` will return a ``np.generic`` if a scalar is passed; not an ``np.ndarray``. --- numpy-stubs/__init__.pyi | 67 +++++++++++++++++++++++++++++++++++++ tests/fail/fromnumeric.py | 37 ++++++++++++++++++++ tests/pass/fromnumeric.py | 57 +++++++++++++++++++++++++++++++ tests/reveal/fromnumeric.py | 57 +++++++++++++++++++++++++++++++ 4 files changed, 218 insertions(+) diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index c999e17..c4d0b9b 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 @@ -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 From c566a8c263c83ad1786c131b00e6fa62e99a6a47 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Fri, 24 Apr 2020 15:28:00 +0200 Subject: [PATCH 2/3] Replace _ArrayLike with _ArrayLikeNested for >= 1D arrays --- numpy-stubs/__init__.pyi | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index c4d0b9b..33fa1d5 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -794,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[ @@ -922,67 +922,67 @@ def argsort( ) -> ndarray: ... @overload def argmax( - a: Union[Sequence[_ArrayLike], ndarray], + a: Union[Sequence[ArrayLike], ndarray], axis: None = ..., out: Optional[ndarray] = ..., ) -> integer: ... @overload def argmax( - a: Union[Sequence[_ArrayLike], ndarray], + a: Union[Sequence[ArrayLike], ndarray], axis: int = ..., out: Optional[ndarray] = ..., ) -> Union[integer, ndarray]: ... @overload def argmin( - a: Union[Sequence[_ArrayLike], ndarray], + a: Union[Sequence[ArrayLike], ndarray], axis: None = ..., out: Optional[ndarray] = ..., ) -> integer: ... @overload def argmin( - a: Union[Sequence[_ArrayLike], ndarray], + a: Union[Sequence[ArrayLike], ndarray], axis: int = ..., out: Optional[ndarray] = ..., ) -> Union[integer, ndarray]: ... @overload def searchsorted( - a: Union[Sequence[_ArrayLike], ndarray], + 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, + 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: ... +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 squeeze(a: ArrayLike, axis: Optional[_ShapeLike] = ...) -> ndarray: ... def diagonal( - a: Union[Sequence[Sequence[_ArrayLike]], ndarray], # >= 2D array + 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 + 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 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, + a: ArrayLike, axis: Optional[int] = ..., out: Optional[ndarray] = ..., ) -> ndarray: ... From aaed219b809ac16e6872c3710abb376ce53d1608 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Wed, 20 May 2020 00:39:48 +0200 Subject: [PATCH 3/3] Changed `_DtypeLike` into `DtypeLike` --- numpy-stubs/__init__.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index 33fa1d5..1d19b6e 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -974,7 +974,7 @@ def trace( offset: int = ..., axis1: int = ..., axis2: int = ..., - dtype: _DtypeLike = ..., + dtype: DtypeLike = ..., out: Optional[ndarray] = ..., ) -> Union[number, ndarray]: ... def ravel(a: ArrayLike, order: _Order = ...) -> ndarray: ...