Merge pull request #17564 from BvB93/numeric

MAINT: Update the annotations in `np.core.numeric`

6 files changed, 351 insertions, 34 deletions

diff --git a/numpy/core/numeric.pyi b/numpy/core/numeric.pyi
index 19720fbdc..f917f74de 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: Union[int, Tuple[_ShapeLike, _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: ...
diff --git a/numpy/typing/tests/data/pass/array_constructors.py b/numpy/typing/tests/data/pass/array_constructors.py
index bf29e52b9..08f4b6211 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 @@ class Index:
 
 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 @@ np.logspace(0, 2, base=2)
 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 000000000..34fef7270
--- /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/pass/simple.py b/numpy/typing/tests/data/pass/simple.py
index 4d397bd29..243caf229 100644
--- a/numpy/typing/tests/data/pass/simple.py
+++ b/numpy/typing/tests/data/pass/simple.py
@@ -17,14 +17,6 @@ ndarray_func(np.array([1, 2]))
 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)
diff --git a/numpy/typing/tests/data/reveal/array_constructors.py b/numpy/typing/tests/data/reveal/array_constructors.py
index 13a04777f..04732ff2e 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 @@ reveal_type(np.linspace(0, 10))  # E: numpy.ndarray
 reveal_type(np.linspace(0, 10, retstep=True))  # E: Tuple[numpy.ndarray, numpy.inexact[Any]]
 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 000000000..6bd4ed9fd
--- /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