Merge pull request #20199 from BvB93/numeric

ENH: Add dtype typing support to `np.core.numeric`

4 files changed, 125 insertions, 77 deletions

diff --git a/numpy/typing/tests/data/fail/array_constructors.pyi b/numpy/typing/tests/data/fail/array_constructors.pyi
index 0e2250513..4f0a60b5b 100644
--- a/numpy/typing/tests/data/fail/array_constructors.pyi
+++ b/numpy/typing/tests/data/fail/array_constructors.pyi
@@ -10,7 +10,7 @@ np.zeros("test")  # E: incompatible type
 np.zeros()  # E: require at least one argument
 
 np.ones("test")  # E: incompatible type
-np.ones()  # E: Missing positional argument
+np.ones()  # E: require at least one argument
 
 np.array(0, float, True)  # E: No overload variant
 
diff --git a/numpy/typing/tests/data/reveal/array_constructors.pyi b/numpy/typing/tests/data/reveal/array_constructors.pyi
index 44c85e988..0aea4ea96 100644
--- a/numpy/typing/tests/data/reveal/array_constructors.pyi
+++ b/numpy/typing/tests/data/reveal/array_constructors.pyi
@@ -120,32 +120,41 @@ reveal_type(np.logspace(0, 10))  # E: numpy.ndarray[Any, Any]
 reveal_type(np.geomspace(1, 10))  # E: numpy.ndarray[Any, Any]
 
 reveal_type(np.zeros_like(A))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
-reveal_type(np.zeros_like(C))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.zeros_like(B))  # E: SubClass
-reveal_type(np.zeros_like(B, dtype=np.int64))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.zeros_like(C))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.zeros_like(A, dtype=float))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.zeros_like(B))  # E: SubClass[{float64}]
+reveal_type(np.zeros_like(B, dtype=np.int64))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
 
 reveal_type(np.ones_like(A))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
-reveal_type(np.ones_like(C))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.ones_like(B))  # E: SubClass
-reveal_type(np.ones_like(B, dtype=np.int64))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.ones_like(C))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.ones_like(A, dtype=float))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.ones_like(B))  # E: SubClass[{float64}]
+reveal_type(np.ones_like(B, dtype=np.int64))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
 
 reveal_type(np.full_like(A, i8))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
-reveal_type(np.full_like(C, i8))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.full_like(C, i8))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.full_like(A, i8, dtype=int))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
 reveal_type(np.full_like(B, i8))  # E: SubClass[{float64}]
-reveal_type(np.full_like(B, i8, dtype=np.int64))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.full_like(B, i8, dtype=np.int64))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
 
-reveal_type(np.ones(1))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.ones([1, 1, 1]))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.ones(1))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.ones([1, 1, 1]))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.ones(5, dtype=np.int64))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
+reveal_type(np.ones(5, dtype=int))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
 
-reveal_type(np.full(1, i8))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.full([1, 1, 1], i8))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.full(1, i8))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.full([1, 1, 1], i8))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.full(1, i8, dtype=np.float64))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.full(1, i8, dtype=float))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
 
-reveal_type(np.indices([1, 2, 3]))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.indices([1, 2, 3], sparse=True))  # E: tuple[numpy.ndarray[Any, Any]]
+reveal_type(np.indices([1, 2, 3]))  # E: numpy.ndarray[Any, numpy.dtype[{int_}]]
+reveal_type(np.indices([1, 2, 3], sparse=True))  # E: tuple[numpy.ndarray[Any, numpy.dtype[{int_}]]]
 
 reveal_type(np.fromfunction(func, (3, 5)))  # E: SubClass[{float64}]
 
-reveal_type(np.identity(10))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.identity(10))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.identity(10, dtype=np.int64))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
+reveal_type(np.identity(10, dtype=int))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
 
 reveal_type(np.atleast_1d(A))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
 reveal_type(np.atleast_1d(C))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
diff --git a/numpy/typing/tests/data/reveal/numeric.pyi b/numpy/typing/tests/data/reveal/numeric.pyi
index ec6e47ca0..9b3b1419d 100644
--- a/numpy/typing/tests/data/reveal/numeric.pyi
+++ b/numpy/typing/tests/data/reveal/numeric.pyi
@@ -7,83 +7,128 @@ Does not include tests which fall under ``array_constructors``.
 
 from typing import List
 import numpy as np
+import numpy.typing as npt
 
-class SubClass(np.ndarray):
+class SubClass(npt.NDArray[np.int64]):
     ...
 
 i8: np.int64
 
-A: np.ndarray
+AR_b: npt.NDArray[np.bool_]
+AR_u8: npt.NDArray[np.uint64]
+AR_i8: npt.NDArray[np.int64]
+AR_f8: npt.NDArray[np.float64]
+AR_c16: npt.NDArray[np.complex128]
+AR_m: npt.NDArray[np.timedelta64]
+AR_O: npt.NDArray[np.object_]
+
 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(AR_i8))  # E: int
 reveal_type(np.count_nonzero(B))  # E: int
-reveal_type(np.count_nonzero(A, keepdims=True))  # E: Any
-reveal_type(np.count_nonzero(A, axis=0))  # E: Any
+reveal_type(np.count_nonzero(AR_i8, keepdims=True))  # E: Any
+reveal_type(np.count_nonzero(AR_i8, axis=0))  # E: Any
 
 reveal_type(np.isfortran(i8))  # E: bool
-reveal_type(np.isfortran(A))  # E: bool
-
-reveal_type(np.argwhere(i8))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.argwhere(A))  # E: numpy.ndarray[Any, Any]
-
-reveal_type(np.flatnonzero(i8))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.flatnonzero(A))  # E: numpy.ndarray[Any, Any]
-
-reveal_type(np.correlate(B, A, mode="valid"))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.correlate(A, A, mode="same"))  # E: numpy.ndarray[Any, Any]
-
-reveal_type(np.convolve(B, A, mode="valid"))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.convolve(A, A, mode="same"))  # E: numpy.ndarray[Any, Any]
-
-reveal_type(np.outer(i8, A))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.outer(B, A))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.outer(A, A))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.outer(A, A, out=C))  # E: SubClass
-
-reveal_type(np.tensordot(B, A))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.tensordot(A, A))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.tensordot(A, A, axes=0))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.tensordot(A, A, axes=(0, 1)))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.isfortran(AR_i8))  # E: bool
+
+reveal_type(np.argwhere(i8))  # E: numpy.ndarray[Any, numpy.dtype[{intp}]]
+reveal_type(np.argwhere(AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[{intp}]]
+
+reveal_type(np.flatnonzero(i8))  # E: numpy.ndarray[Any, numpy.dtype[{intp}]]
+reveal_type(np.flatnonzero(AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[{intp}]]
+
+reveal_type(np.correlate(B, AR_i8, mode="valid"))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.correlate(AR_i8, AR_i8, mode="same"))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.correlate(AR_b, AR_b))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+reveal_type(np.correlate(AR_b, AR_u8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.unsignedinteger[Any]]]
+reveal_type(np.correlate(AR_i8, AR_b))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.correlate(AR_i8, AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
+reveal_type(np.correlate(AR_i8, AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+reveal_type(np.correlate(AR_i8, AR_m))  # E: numpy.ndarray[Any, numpy.dtype[numpy.timedelta64]]
+reveal_type(np.correlate(AR_O, AR_O))  # E: numpy.ndarray[Any, numpy.dtype[numpy.object_]]
+
+reveal_type(np.convolve(B, AR_i8, mode="valid"))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.convolve(AR_i8, AR_i8, mode="same"))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.convolve(AR_b, AR_b))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+reveal_type(np.convolve(AR_b, AR_u8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.unsignedinteger[Any]]]
+reveal_type(np.convolve(AR_i8, AR_b))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.convolve(AR_i8, AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
+reveal_type(np.convolve(AR_i8, AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+reveal_type(np.convolve(AR_i8, AR_m))  # E: numpy.ndarray[Any, numpy.dtype[numpy.timedelta64]]
+reveal_type(np.convolve(AR_O, AR_O))  # E: numpy.ndarray[Any, numpy.dtype[numpy.object_]]
+
+reveal_type(np.outer(i8, AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.outer(B, AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.outer(AR_i8, AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.outer(AR_i8, AR_i8, out=C))  # E: SubClass
+reveal_type(np.outer(AR_b, AR_b))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+reveal_type(np.outer(AR_b, AR_u8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.unsignedinteger[Any]]]
+reveal_type(np.outer(AR_i8, AR_b))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.convolve(AR_i8, AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
+reveal_type(np.outer(AR_i8, AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+reveal_type(np.outer(AR_i8, AR_m))  # E: numpy.ndarray[Any, numpy.dtype[numpy.timedelta64]]
+reveal_type(np.outer(AR_O, AR_O))  # E: numpy.ndarray[Any, numpy.dtype[numpy.object_]]
+
+reveal_type(np.tensordot(B, AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.tensordot(AR_i8, AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.tensordot(AR_i8, AR_i8, axes=0))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.tensordot(AR_i8, AR_i8, axes=(0, 1)))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.tensordot(AR_b, AR_b))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+reveal_type(np.tensordot(AR_b, AR_u8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.unsignedinteger[Any]]]
+reveal_type(np.tensordot(AR_i8, AR_b))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.tensordot(AR_i8, AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
+reveal_type(np.tensordot(AR_i8, AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+reveal_type(np.tensordot(AR_i8, AR_m))  # E: numpy.ndarray[Any, numpy.dtype[numpy.timedelta64]]
+reveal_type(np.tensordot(AR_O, AR_O))  # E: numpy.ndarray[Any, numpy.dtype[numpy.object_]]
 
 reveal_type(np.isscalar(i8))  # E: bool
-reveal_type(np.isscalar(A))  # E: bool
+reveal_type(np.isscalar(AR_i8))  # E: bool
 reveal_type(np.isscalar(B))  # E: bool
 
-reveal_type(np.roll(A, 1))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.roll(A, (1, 2)))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.roll(B, 1))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.roll(AR_i8, 1))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
+reveal_type(np.roll(AR_i8, (1, 2)))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
+reveal_type(np.roll(B, 1))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
 
-reveal_type(np.rollaxis(A, 0, 1))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.rollaxis(AR_i8, 0, 1))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
 
-reveal_type(np.moveaxis(A, 0, 1))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.moveaxis(A, (0, 1), (1, 2)))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.moveaxis(AR_i8, 0, 1))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
+reveal_type(np.moveaxis(AR_i8, (0, 1), (1, 2)))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
 
-reveal_type(np.cross(B, A))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.cross(A, A))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.cross(B, AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.cross(AR_i8, AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.cross(AR_b, AR_u8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.unsignedinteger[Any]]]
+reveal_type(np.cross(AR_i8, AR_b))  # E: numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+reveal_type(np.cross(AR_i8, AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
+reveal_type(np.cross(AR_i8, AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+reveal_type(np.cross(AR_O, AR_O))  # E: numpy.ndarray[Any, numpy.dtype[numpy.object_]]
 
-reveal_type(np.indices([0, 1, 2]))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.indices([0, 1, 2], sparse=False))  # E: numpy.ndarray[Any, Any]
-reveal_type(np.indices([0, 1, 2], sparse=True))  # E: tuple[numpy.ndarray[Any, Any]]
+reveal_type(np.indices([0, 1, 2]))  # E: numpy.ndarray[Any, numpy.dtype[{int_}]]
+reveal_type(np.indices([0, 1, 2], sparse=True))  # E: tuple[numpy.ndarray[Any, numpy.dtype[{int_}]]]
+reveal_type(np.indices([0, 1, 2], dtype=np.float64))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.indices([0, 1, 2], sparse=True, dtype=np.float64))  # E: tuple[numpy.ndarray[Any, numpy.dtype[{float64}]]]
+reveal_type(np.indices([0, 1, 2], dtype=float))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.indices([0, 1, 2], sparse=True, dtype=float))  # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
 
 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.allclose(i8, AR_i8))  # E: bool
+reveal_type(np.allclose(B, AR_i8))  # E: bool
+reveal_type(np.allclose(AR_i8, AR_i8))  # E: bool
 
-reveal_type(np.isclose(i8, A))  # E: Any
-reveal_type(np.isclose(B, A))  # E: Any
-reveal_type(np.isclose(A, A))  # E: Any
+reveal_type(np.isclose(i8, i8))  # E: numpy.bool_
+reveal_type(np.isclose(i8, AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+reveal_type(np.isclose(B, AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+reveal_type(np.isclose(AR_i8, AR_i8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
 
-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_equal(i8, AR_i8))  # E: bool
+reveal_type(np.array_equal(B, AR_i8))  # E: bool
+reveal_type(np.array_equal(AR_i8, AR_i8))  # 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
+reveal_type(np.array_equiv(i8, AR_i8))  # E: bool
+reveal_type(np.array_equiv(B, AR_i8))  # E: bool
+reveal_type(np.array_equiv(AR_i8, AR_i8))  # E: bool
diff --git a/numpy/typing/tests/test_typing.py b/numpy/typing/tests/test_typing.py
index f303ebea3..4dd6530ff 100644
--- a/numpy/typing/tests/test_typing.py
+++ b/numpy/typing/tests/test_typing.py
@@ -98,16 +98,10 @@ def run_mypy() -> None:
 def get_test_cases(directory: str) -> Iterator[ParameterSet]:
     for root, _, files in os.walk(directory):
         for fname in files:
-            if os.path.splitext(fname)[-1] in (".pyi", ".py"):
+            short_fname, ext = os.path.splitext(fname)
+            if ext in (".pyi", ".py"):
                 fullpath = os.path.join(root, fname)
-                # Use relative path for nice py.test name
-                relpath = os.path.relpath(fullpath, start=directory)
-
-                yield pytest.param(
-                    fullpath,
-                    # Manually specify a name for the test
-                    id=relpath,
-                )
+                yield pytest.param(fullpath, id=short_fname)
 
 
 @pytest.mark.slow