From 5b028ae4ec0d23e3fe2e993bd0d1639c48c7b492 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Wed, 14 Oct 2020 15:56:38 +0200 Subject: [PATCH 1/4] MAINT: Updated the annotations in `np.core.numeric` * Added missing overloads * Use `DtypeLike` and `_Shapelike` where applicable * Moved a few annotation-aliases to the top of the module --- numpy/core/numeric.pyi | 120 ++++++++++++++++++++++++++++++++--------- 1 file changed, 96 insertions(+), 24 deletions(-) diff --git a/numpy/core/numeric.pyi b/numpy/core/numeric.pyi index 19720fbdcc9b..b341d136a10f 100644 --- a/numpy/core/numeric.pyi +++ b/numpy/core/numeric.pyi @@ -1,14 +1,45 @@ -from typing import Any, Optional, Union, Sequence, Tuple +import sys +from typing import ( + Any, + Optional, + Union, + Sequence, + Tuple, + Callable, + List, + overload, + TypeVar, + Iterable, +) -from numpy import ndarray, dtype, bool_, _OrderKACF, _OrderCF +from numpy import ndarray, generic, dtype, bool_, int32, int64, _OrderKACF, _OrderCF from numpy.typing import ArrayLike, DtypeLike, _ShapeLike +if sys.version_info >= (3, 8): + from typing import Literal +else: + from typing_extensions import Literal + +_T = TypeVar("_T") +_ArrayType = TypeVar("_ArrayType", bound=ndarray) + +_CorrelateMode = Literal["valid", "same", "full"] + +@overload +def zeros_like( + a: _ArrayType, + dtype: None = ..., + order: _OrderKACF = ..., + subok: Literal[True] = ..., + shape: None = ..., +) -> _ArrayType: ... +@overload def zeros_like( a: ArrayLike, dtype: DtypeLike = ..., order: _OrderKACF = ..., subok: bool = ..., - shape: Optional[Union[int, Sequence[int]]] = ..., + shape: Optional[_ShapeLike] = ..., ) -> ndarray: ... def ones( shape: _ShapeLike, @@ -17,6 +48,15 @@ def ones( *, like: ArrayLike = ..., ) -> ndarray: ... +@overload +def ones_like( + a: _ArrayType, + dtype: None = ..., + order: _OrderKACF = ..., + subok: Literal[True] = ..., + shape: None = ..., +) -> _ArrayType: ... +@overload def ones_like( a: ArrayLike, dtype: DtypeLike = ..., @@ -24,6 +64,15 @@ def ones_like( subok: bool = ..., shape: Optional[_ShapeLike] = ..., ) -> ndarray: ... +@overload +def empty_like( + a: _ArrayType, + dtype: None = ..., + order: _OrderKACF = ..., + subok: Literal[True] = ..., + shape: None = ..., +) -> _ArrayType: ... +@overload def empty_like( a: ArrayLike, dtype: DtypeLike = ..., @@ -39,6 +88,16 @@ def full( *, like: ArrayLike = ..., ) -> ndarray: ... +@overload +def full_like( + a: _ArrayType, + fill_value: Any, + dtype: None = ..., + order: _OrderKACF = ..., + subok: Literal[True] = ..., + shape: None = ..., +) -> _ArrayType: ... +@overload def full_like( a: ArrayLike, fill_value: Any, @@ -47,35 +106,38 @@ def full_like( subok: bool = ..., shape: Optional[_ShapeLike] = ..., ) -> ndarray: ... +@overload +def count_nonzero( + a: ArrayLike, axis: None = ..., *, keepdims: Literal[False] = ... +) -> int: ... +@overload def count_nonzero( - a: ArrayLike, axis: Optional[Union[int, Tuple[int], Tuple[int, int]]] = ... -) -> Union[int, ndarray]: ... -def isfortran(a: ndarray) -> bool: ... + a: ArrayLike, axis: _ShapeLike = ..., *, keepdims: bool = ... +) -> Union[int64, int32, ndarray]: ... +def isfortran(a: Union[ndarray, generic]) -> bool: ... def argwhere(a: ArrayLike) -> ndarray: ... def flatnonzero(a: ArrayLike) -> ndarray: ... - -_CorrelateMode = Literal["valid", "same", "full"] - def correlate(a: ArrayLike, v: ArrayLike, mode: _CorrelateMode = ...) -> ndarray: ... def convolve(a: ArrayLike, v: ArrayLike, mode: _CorrelateMode = ...) -> ndarray: ... -def outer(a: ArrayLike, b: ArrayLike, out: ndarray = ...) -> ndarray: ... +@overload +def outer(a: ArrayLike, b: ArrayLike, out: None = ...) -> ndarray: ... +@overload +def outer(a: ArrayLike, b: ArrayLike, out: _ArrayType = ...) -> _ArrayType: ... def tensordot( a: ArrayLike, b: ArrayLike, - axes: Union[ - int, Tuple[int, int], Tuple[Tuple[int, int], ...], Tuple[List[int, int], ...] - ] = ..., + axes: _ShapeLike = ..., ) -> ndarray: ... def roll( a: ArrayLike, - shift: Union[int, Tuple[int, ...]], - axis: Optional[Union[int, Tuple[int, ...]]] = ..., + shift: _ShapeLike, + axis: Optional[_ShapeLike] = ..., ) -> ndarray: ... -def rollaxis(a: ArrayLike, axis: int, start: int = ...) -> ndarray: ... +def rollaxis(a: ndarray, axis: int, start: int = ...) -> ndarray: ... def moveaxis( a: ndarray, - source: Union[int, Sequence[int]], - destination: Union[int, Sequence[int]], + source: _ShapeLike, + destination: _ShapeLike, ) -> ndarray: ... def cross( a: ArrayLike, @@ -85,16 +147,26 @@ def cross( axisc: int = ..., axis: Optional[int] = ..., ) -> ndarray: ... +@overload def indices( - dimensions: Sequence[int], dtype: dtype = ..., sparse: bool = ... -) -> Union[ndarray, Tuple[ndarray, ...]]: ... + dimensions: Sequence[int], + dtype: DtypeLike = ..., + sparse: Literal[False] = ..., +) -> ndarray: ... +@overload +def indices( + dimensions: Sequence[int], + dtype: DtypeLike = ..., + sparse: Literal[True] = ..., +) -> Tuple[ndarray, ...]: ... def fromfunction( - function: Callable, - shape: Tuple[int, int], + function: Callable[..., _T], + shape: Sequence[int], *, + dtype: DtypeLike = ..., like: ArrayLike = ..., - **kwargs, -) -> Any: ... + **kwargs: Any, +) -> _T: ... def isscalar(element: Any) -> bool: ... def binary_repr(num: int, width: Optional[int] = ...) -> str: ... def base_repr(number: int, base: int = ..., padding: int = ...) -> str: ... From 4e3c6bd71b170667c8c73f0d28fe6f67f5a4252b Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Wed, 14 Oct 2020 15:57:03 +0200 Subject: [PATCH 2/4] TST,MAINT: Moved a few tests to `array_constructors` --- numpy/typing/tests/data/pass/simple.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/numpy/typing/tests/data/pass/simple.py b/numpy/typing/tests/data/pass/simple.py index 4d397bd29f29..243caf229f13 100644 --- a/numpy/typing/tests/data/pass/simple.py +++ b/numpy/typing/tests/data/pass/simple.py @@ -17,14 +17,6 @@ def ndarray_func(x): array == 1 array.dtype == float -ndarray_func(np.zeros([1, 2])) -ndarray_func(np.ones([1, 2])) -ndarray_func(np.empty([1, 2])) - -ndarray_func(np.zeros_like(array)) -ndarray_func(np.ones_like(array)) -ndarray_func(np.empty_like(array)) - # Dtype construction np.dtype(float) np.dtype(np.float64) From 71292e07e0329d0b256804146d06b99d83fd36c9 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Wed, 14 Oct 2020 16:57:08 +0200 Subject: [PATCH 3/4] TST: Added tests for `np.core.numeric` --- .../tests/data/pass/array_constructors.py | 40 ++++++++- numpy/typing/tests/data/pass/numeric.py | 89 +++++++++++++++++++ .../tests/data/reveal/array_constructors.py | 39 +++++++- numpy/typing/tests/data/reveal/numeric.py | 89 +++++++++++++++++++ 4 files changed, 255 insertions(+), 2 deletions(-) create mode 100644 numpy/typing/tests/data/pass/numeric.py create mode 100644 numpy/typing/tests/data/reveal/numeric.py diff --git a/numpy/typing/tests/data/pass/array_constructors.py b/numpy/typing/tests/data/pass/array_constructors.py index bf29e52b91b9..08f4b6211dc9 100644 --- a/numpy/typing/tests/data/pass/array_constructors.py +++ b/numpy/typing/tests/data/pass/array_constructors.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Any import numpy as np class Index: @@ -7,10 +7,15 @@ def __index__(self) -> int: class SubClass(np.ndarray): ... +i8 = np.int64(1) + A = np.array([1]) B = A.view(SubClass).copy() C = [1] +def func(i: int, j: int, **kwargs: Any) -> SubClass: + return B + np.array(1, dtype=float) np.array(1, copy=False) np.array(1, order='F') @@ -64,3 +69,36 @@ class SubClass(np.ndarray): ... np.logspace(0, 2, base=[1j, 2j], num=2) np.geomspace(1, 2) + +np.zeros_like(A) +np.zeros_like(C) +np.zeros_like(B) +np.zeros_like(B, dtype=np.int64) + +np.ones_like(A) +np.ones_like(C) +np.ones_like(B) +np.ones_like(B, dtype=np.int64) + +np.empty_like(A) +np.empty_like(C) +np.empty_like(B) +np.empty_like(B, dtype=np.int64) + +np.full_like(A, i8) +np.full_like(C, i8) +np.full_like(B, i8) +np.full_like(B, i8, dtype=np.int64) + +np.ones(1) +np.ones([1, 1, 1]) + +np.full(1, i8) +np.full([1, 1, 1], i8) + +np.indices([1, 2, 3]) +np.indices([1, 2, 3], sparse=True) + +np.fromfunction(func, (3, 5)) + +np.identity(10) diff --git a/numpy/typing/tests/data/pass/numeric.py b/numpy/typing/tests/data/pass/numeric.py new file mode 100644 index 000000000000..34fef7270443 --- /dev/null +++ b/numpy/typing/tests/data/pass/numeric.py @@ -0,0 +1,89 @@ +""" +Tests for :mod:`numpy.core.numeric`. + +Does not include tests which fall under ``array_constructors``. + +""" + +from typing import List +import numpy as np + +class SubClass(np.ndarray): + ... + +i8 = np.int64(1) + +A = np.arange(27).reshape(3, 3, 3) +B: List[List[List[int]]] = A.tolist() +C = np.empty((27, 27)).view(SubClass) + +np.count_nonzero(i8) +np.count_nonzero(A) +np.count_nonzero(B) +np.count_nonzero(A, keepdims=True) +np.count_nonzero(A, axis=0) + +np.isfortran(i8) +np.isfortran(A) + +np.argwhere(i8) +np.argwhere(A) + +np.flatnonzero(i8) +np.flatnonzero(A) + +np.correlate(B[0][0], A.ravel(), mode="valid") +np.correlate(A.ravel(), A.ravel(), mode="same") + +np.convolve(B[0][0], A.ravel(), mode="valid") +np.convolve(A.ravel(), A.ravel(), mode="same") + +np.outer(i8, A) +np.outer(B, A) +np.outer(A, A) +np.outer(A, A, out=C) + +np.tensordot(B, A) +np.tensordot(A, A) +np.tensordot(A, A, axes=0) +np.tensordot(A, A, axes=(0, 1)) + +np.isscalar(i8) +np.isscalar(A) +np.isscalar(B) + +np.roll(A, 1) +np.roll(A, (1, 2)) +np.roll(B, 1) + +np.rollaxis(A, 0, 1) + +np.moveaxis(A, 0, 1) +np.moveaxis(A, (0, 1), (1, 2)) + +np.cross(B, A) +np.cross(A, A) + +np.indices([0, 1, 2]) +np.indices([0, 1, 2], sparse=False) +np.indices([0, 1, 2], sparse=True) + +np.binary_repr(1) + +np.base_repr(1) + +np.allclose(i8, A) +np.allclose(B, A) +np.allclose(A, A) + +np.isclose(i8, A) +np.isclose(B, A) +np.isclose(A, A) + +np.array_equal(i8, A) +np.array_equal(B, A) +np.array_equal(A, A) + +np.array_equiv(i8, A) +np.array_equiv(B, A) +np.array_equiv(A, A) diff --git a/numpy/typing/tests/data/reveal/array_constructors.py b/numpy/typing/tests/data/reveal/array_constructors.py index ba8a8eda1e94..9475b3884b9a 100644 --- a/numpy/typing/tests/data/reveal/array_constructors.py +++ b/numpy/typing/tests/data/reveal/array_constructors.py @@ -1,12 +1,16 @@ -from typing import List +from typing import List, Any import numpy as np class SubClass(np.ndarray): ... +i8: np.int64 + A: np.ndarray B: SubClass C: List[int] +def func(i: int, j: int, **kwargs: Any) -> SubClass: ... + reveal_type(np.asarray(A)) # E: ndarray reveal_type(np.asarray(B)) # E: ndarray reveal_type(np.asarray(C)) # E: ndarray @@ -40,3 +44,36 @@ class SubClass(np.ndarray): ... reveal_type(np.linspace(0, 10, retstep=True)) # E: Tuple[numpy.ndarray, numpy.inexact] reveal_type(np.logspace(0, 10)) # E: numpy.ndarray reveal_type(np.geomspace(1, 10)) # E: numpy.ndarray + +reveal_type(np.zeros_like(A)) # E: numpy.ndarray +reveal_type(np.zeros_like(C)) # E: numpy.ndarray +reveal_type(np.zeros_like(B)) # E: SubClass +reveal_type(np.zeros_like(B, dtype=np.int64)) # E: numpy.ndarray + +reveal_type(np.ones_like(A)) # E: numpy.ndarray +reveal_type(np.ones_like(C)) # E: numpy.ndarray +reveal_type(np.ones_like(B)) # E: SubClass +reveal_type(np.ones_like(B, dtype=np.int64)) # E: numpy.ndarray + +reveal_type(np.empty_like(A)) # E: numpy.ndarray +reveal_type(np.empty_like(C)) # E: numpy.ndarray +reveal_type(np.empty_like(B)) # E: SubClass +reveal_type(np.empty_like(B, dtype=np.int64)) # E: numpy.ndarray + +reveal_type(np.full_like(A, i8)) # E: numpy.ndarray +reveal_type(np.full_like(C, i8)) # E: numpy.ndarray +reveal_type(np.full_like(B, i8)) # E: SubClass +reveal_type(np.full_like(B, i8, dtype=np.int64)) # E: numpy.ndarray + +reveal_type(np.ones(1)) # E: numpy.ndarray +reveal_type(np.ones([1, 1, 1])) # E: numpy.ndarray + +reveal_type(np.full(1, i8)) # E: numpy.ndarray +reveal_type(np.full([1, 1, 1], i8)) # E: numpy.ndarray + +reveal_type(np.indices([1, 2, 3])) # E: numpy.ndarray +reveal_type(np.indices([1, 2, 3], sparse=True)) # E: tuple[numpy.ndarray] + +reveal_type(np.fromfunction(func, (3, 5))) # E: SubClass + +reveal_type(np.identity(10)) # E: numpy.ndarray diff --git a/numpy/typing/tests/data/reveal/numeric.py b/numpy/typing/tests/data/reveal/numeric.py new file mode 100644 index 000000000000..6bd4ed9fdd1e --- /dev/null +++ b/numpy/typing/tests/data/reveal/numeric.py @@ -0,0 +1,89 @@ +""" +Tests for :mod:`numpy.core.numeric`. + +Does not include tests which fall under ``array_constructors``. + +""" + +from typing import List +import numpy as np + +class SubClass(np.ndarray): + ... + +i8: np.int64 + +A: np.ndarray +B: List[int] +C: SubClass + +reveal_type(np.count_nonzero(i8)) # E: int +reveal_type(np.count_nonzero(A)) # E: int +reveal_type(np.count_nonzero(B)) # E: int +reveal_type(np.count_nonzero(A, keepdims=True)) # E: Union[numpy.int64, numpy.int32, numpy.ndarray] +reveal_type(np.count_nonzero(A, axis=0)) # E: Union[numpy.int64, numpy.int32, numpy.ndarray] + +reveal_type(np.isfortran(i8)) # E: bool +reveal_type(np.isfortran(A)) # E: bool + +reveal_type(np.argwhere(i8)) # E: numpy.ndarray +reveal_type(np.argwhere(A)) # E: numpy.ndarray + +reveal_type(np.flatnonzero(i8)) # E: numpy.ndarray +reveal_type(np.flatnonzero(A)) # E: numpy.ndarray + +reveal_type(np.correlate(B, A, mode="valid")) # E: numpy.ndarray +reveal_type(np.correlate(A, A, mode="same")) # E: numpy.ndarray + +reveal_type(np.convolve(B, A, mode="valid")) # E: numpy.ndarray +reveal_type(np.convolve(A, A, mode="same")) # E: numpy.ndarray + +reveal_type(np.outer(i8, A)) # E: numpy.ndarray +reveal_type(np.outer(B, A)) # E: numpy.ndarray +reveal_type(np.outer(A, A)) # E: numpy.ndarray +reveal_type(np.outer(A, A, out=C)) # E: SubClass + +reveal_type(np.tensordot(B, A)) # E: numpy.ndarray +reveal_type(np.tensordot(A, A)) # E: numpy.ndarray +reveal_type(np.tensordot(A, A, axes=0)) # E: numpy.ndarray +reveal_type(np.tensordot(A, A, axes=(0, 1))) # E: numpy.ndarray + +reveal_type(np.isscalar(i8)) # E: bool +reveal_type(np.isscalar(A)) # E: bool +reveal_type(np.isscalar(B)) # E: bool + +reveal_type(np.roll(A, 1)) # E: numpy.ndarray +reveal_type(np.roll(A, (1, 2))) # E: numpy.ndarray +reveal_type(np.roll(B, 1)) # E: numpy.ndarray + +reveal_type(np.rollaxis(A, 0, 1)) # E: numpy.ndarray + +reveal_type(np.moveaxis(A, 0, 1)) # E: numpy.ndarray +reveal_type(np.moveaxis(A, (0, 1), (1, 2))) # E: numpy.ndarray + +reveal_type(np.cross(B, A)) # E: numpy.ndarray +reveal_type(np.cross(A, A)) # E: numpy.ndarray + +reveal_type(np.indices([0, 1, 2])) # E: numpy.ndarray +reveal_type(np.indices([0, 1, 2], sparse=False)) # E: numpy.ndarray +reveal_type(np.indices([0, 1, 2], sparse=True)) # E: tuple[numpy.ndarray] + +reveal_type(np.binary_repr(1)) # E: str + +reveal_type(np.base_repr(1)) # E: str + +reveal_type(np.allclose(i8, A)) # E: bool +reveal_type(np.allclose(B, A)) # E: bool +reveal_type(np.allclose(A, A)) # E: bool + +reveal_type(np.isclose(i8, A)) # E: Union[numpy.bool_, numpy.ndarray] +reveal_type(np.isclose(B, A)) # E: Union[numpy.bool_, numpy.ndarray] +reveal_type(np.isclose(A, A)) # E: Union[numpy.bool_, numpy.ndarray] + +reveal_type(np.array_equal(i8, A)) # E: bool +reveal_type(np.array_equal(B, A)) # E: bool +reveal_type(np.array_equal(A, A)) # E: bool + +reveal_type(np.array_equiv(i8, A)) # E: bool +reveal_type(np.array_equiv(B, A)) # E: bool +reveal_type(np.array_equiv(A, A)) # E: bool From 3a6857c99d9012650b8b58726783eee30308d944 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Wed, 21 Oct 2020 13:01:11 +0200 Subject: [PATCH 4/4] MAINT: Update the `axes` parameter of `tensordot` `axes` can take either an integer or a 2-tuple of shape-like objects --- numpy/core/numeric.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/numpy/core/numeric.pyi b/numpy/core/numeric.pyi index b341d136a10f..f917f74dec41 100644 --- a/numpy/core/numeric.pyi +++ b/numpy/core/numeric.pyi @@ -126,7 +126,7 @@ def outer(a: ArrayLike, b: ArrayLike, out: _ArrayType = ...) -> _ArrayType: ... def tensordot( a: ArrayLike, b: ArrayLike, - axes: _ShapeLike = ..., + axes: Union[int, Tuple[_ShapeLike, _ShapeLike]] = ..., ) -> ndarray: ... def roll( a: ArrayLike,