TYP,ENH: Add dtype-typing support to `fromnumeric` (part 1)

2 files changed, 89 insertions, 85 deletions

diff --git a/numpy/typing/tests/data/fail/fromnumeric.pyi b/numpy/typing/tests/data/fail/fromnumeric.pyi
index 8fafed1b7..84c95639d 100644
--- a/numpy/typing/tests/data/fail/fromnumeric.pyi
+++ b/numpy/typing/tests/data/fail/fromnumeric.pyi
@@ -12,27 +12,27 @@ np.take(a, axis=1.0)  # E: incompatible type
 np.take(A, out=1)  # E: incompatible type
 np.take(A, mode="bob")  # E: incompatible type
 
-np.reshape(a, None)  # E: Argument 2 to "reshape" has incompatible type
-np.reshape(A, 1, order="bob")  # E: Argument "order" to "reshape" has incompatible type
+np.reshape(a, None)  # E: No overload variant
+np.reshape(A, 1, order="bob")  # E: No overload variant
 
-np.choose(a, None)  # E: incompatible type
-np.choose(a, out=1.0)  # E: incompatible type
-np.choose(A, mode="bob")  # E: incompatible type
+np.choose(a, None)  # E: No overload variant
+np.choose(a, out=1.0)  # E: No overload variant
+np.choose(A, mode="bob")  # E: No overload variant
 
-np.repeat(a, None)  # E: Argument 2 to "repeat" has incompatible type
-np.repeat(A, 1, axis=1.0)  # E: Argument "axis" to "repeat" has incompatible type
+np.repeat(a, None)  # E: No overload variant
+np.repeat(A, 1, axis=1.0)  # E: No overload variant
 
-np.swapaxes(A, None, 1)  # E: Argument 2 to "swapaxes" has incompatible type
-np.swapaxes(A, 1, [0])  # E: Argument 3 to "swapaxes" has incompatible type
+np.swapaxes(A, None, 1)  # E: No overload variant
+np.swapaxes(A, 1, [0])  # E: No overload variant
 
-np.transpose(A, axes=1.0)  # E: Argument "axes" to "transpose" has incompatible type
+np.transpose(A, axes=1.0)  # E: No overload variant
 
-np.partition(a, None)  # E: Argument 2 to "partition" has incompatible type
-np.partition(
-    a, 0, axis="bob"  # E: Argument "axis" to "partition" has incompatible type
+np.partition(a, None)  # E: No overload variant
+np.partition(  # E: No overload variant
+    a, 0, axis="bob"
 )
-np.partition(
-    A, 0, kind="bob"  # E: Argument "kind" to "partition" has incompatible type
+np.partition(  # E: No overload variant
+    A, 0, kind="bob"
 )
 np.partition(
     A, 0, order=range(5)  # E: Argument "order" to "partition" has incompatible type
@@ -51,8 +51,8 @@ np.argpartition(
     A, 0, order=range(5)  # E: Argument "order" to "argpartition" has incompatible type
 )
 
-np.sort(A, axis="bob")  # E: Argument "axis" to "sort" has incompatible type
-np.sort(A, kind="bob")  # E: Argument "kind" to "sort" has incompatible type
+np.sort(A, axis="bob")  # E: No overload variant
+np.sort(A, kind="bob")  # E: No overload variant
 np.sort(A, order=range(5))  # E: Argument "order" to "sort" has incompatible type
 
 np.argsort(A, axis="bob")  # E: Argument "axis" to "argsort" has incompatible type
@@ -72,22 +72,22 @@ np.searchsorted(  # E: No overload variant of "searchsorted" matches argument ty
     A[0], 0, sorter=1.0
 )
 
-np.resize(A, 1.0)  # E: Argument 2 to "resize" has incompatible type
+np.resize(A, 1.0)  # E: No overload variant
 
 np.squeeze(A, 1.0)  # E: No overload variant of "squeeze" matches argument 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.diagonal(A, offset=None)  # E: No overload variant
+np.diagonal(A, axis1="bob")  # E: No overload variant
+np.diagonal(A, axis2=[])  # E: No overload variant
 
 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.ravel(a, order="bob")  # E: No overload variant
 
-np.compress(
-    [True], A, axis=1.0  # E: Argument "axis" to "compress" has incompatible type
+np.compress(  # E: No overload variant
+    [True], A, axis=1.0
 )
 
 np.clip(a, 1, 2, out=1)  # E: No overload variant of "clip" matches argument type
diff --git a/numpy/typing/tests/data/reveal/fromnumeric.pyi b/numpy/typing/tests/data/reveal/fromnumeric.pyi
index 2ee1952cf..aba570f1a 100644
--- a/numpy/typing/tests/data/reveal/fromnumeric.pyi
+++ b/numpy/typing/tests/data/reveal/fromnumeric.pyi
@@ -2,7 +2,7 @@
 
 import numpy as np
 
-A = np.array(True, ndmin=2, dtype=bool)
+A = np.array(True, ndmin=2, dtype=np.bool_)
 B = np.array(1.0, ndmin=2, dtype=np.float32)
 A.setflags(write=False)
 B.setflags(write=False)
@@ -20,49 +20,53 @@ reveal_type(np.take(B, 0))  # E: Any
 reveal_type(np.take(A, [0]))  # E: Any
 reveal_type(np.take(B, [0]))  # E: Any
 
-reveal_type(np.reshape(a, 1))  # E: ndarray[Any, Any]
-reveal_type(np.reshape(b, 1))  # E: ndarray[Any, Any]
-reveal_type(np.reshape(c, 1))  # E: ndarray[Any, Any]
-reveal_type(np.reshape(A, 1))  # E: ndarray[Any, Any]
-reveal_type(np.reshape(B, 1))  # E: ndarray[Any, Any]
+reveal_type(np.reshape(a, 1))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.reshape(b, 1))  # E: ndarray[Any, dtype[{float32}]]
+reveal_type(np.reshape(c, 1))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.reshape(A, 1))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.reshape(B, 1))  # E: ndarray[Any, dtype[{float32}]]
 
-reveal_type(np.choose(a, [True, True]))  # E: Any
-reveal_type(np.choose(A, [True, True]))  # E: Any
+reveal_type(np.choose(1, [True, True]))  # E: Any
+reveal_type(np.choose([1], [True, True]))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.choose([1], A))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.choose([1], A, out=B))  # E: ndarray[Any, dtype[{float32}]]
 
-reveal_type(np.repeat(a, 1))  # E: ndarray[Any, Any]
-reveal_type(np.repeat(b, 1))  # E: ndarray[Any, Any]
-reveal_type(np.repeat(c, 1))  # E: ndarray[Any, Any]
-reveal_type(np.repeat(A, 1))  # E: ndarray[Any, Any]
-reveal_type(np.repeat(B, 1))  # E: ndarray[Any, Any]
+reveal_type(np.repeat(a, 1))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.repeat(b, 1))  # E: ndarray[Any, dtype[{float32}]]
+reveal_type(np.repeat(c, 1))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.repeat(A, 1))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.repeat(B, 1))  # E: ndarray[Any, dtype[{float32}]]
 
 # TODO: Add tests for np.put()
 
-reveal_type(np.swapaxes(A, 0, 0))  # E: ndarray[Any, Any]
-reveal_type(np.swapaxes(B, 0, 0))  # E: ndarray[Any, Any]
+reveal_type(np.swapaxes([[0, 1]], 0, 0))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.swapaxes(A, 0, 0))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.swapaxes(B, 0, 0))  # E: ndarray[Any, dtype[{float32}]]
 
-reveal_type(np.transpose(a))  # E: ndarray[Any, Any]
-reveal_type(np.transpose(b))  # E: ndarray[Any, Any]
-reveal_type(np.transpose(c))  # E: ndarray[Any, Any]
-reveal_type(np.transpose(A))  # E: ndarray[Any, Any]
-reveal_type(np.transpose(B))  # E: ndarray[Any, Any]
+reveal_type(np.transpose(a))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.transpose(b))  # E: ndarray[Any, dtype[{float32}]]
+reveal_type(np.transpose(c))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.transpose(A))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.transpose(B))  # E: ndarray[Any, dtype[{float32}]]
 
-reveal_type(np.partition(a, 0, axis=None))  # E: ndarray[Any, Any]
-reveal_type(np.partition(b, 0, axis=None))  # E: ndarray[Any, Any]
-reveal_type(np.partition(c, 0, axis=None))  # E: ndarray[Any, Any]
-reveal_type(np.partition(A, 0))  # E: ndarray[Any, Any]
-reveal_type(np.partition(B, 0))  # E: ndarray[Any, Any]
+reveal_type(np.partition(a, 0, axis=None))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.partition(b, 0, axis=None))  # E: ndarray[Any, dtype[{float32}]]
+reveal_type(np.partition(c, 0, axis=None))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.partition(A, 0))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.partition(B, 0))  # E: ndarray[Any, dtype[{float32}]]
 
-reveal_type(np.argpartition(a, 0))  # E: Any
-reveal_type(np.argpartition(b, 0))  # E: Any
-reveal_type(np.argpartition(c, 0))  # E: Any
-reveal_type(np.argpartition(A, 0))  # E: Any
-reveal_type(np.argpartition(B, 0))  # E: Any
+reveal_type(np.argpartition(a, 0))  # E: ndarray[Any, dtype[{intp}]]
+reveal_type(np.argpartition(b, 0))  # E: ndarray[Any, dtype[{intp}]]
+reveal_type(np.argpartition(c, 0))  # E: ndarray[Any, dtype[{intp}]]
+reveal_type(np.argpartition(A, 0))  # E: ndarray[Any, dtype[{intp}]]
+reveal_type(np.argpartition(B, 0))  # E: ndarray[Any, dtype[{intp}]]
 
-reveal_type(np.sort(A, 0))  # E: ndarray[Any, Any]
-reveal_type(np.sort(B, 0))  # E: ndarray[Any, Any]
+reveal_type(np.sort([2, 1], 0))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.sort(A, 0))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.sort(B, 0))  # E: ndarray[Any, dtype[{float32}]]
 
-reveal_type(np.argsort(A, 0))  # E: ndarray[Any, Any]
-reveal_type(np.argsort(B, 0))  # E: ndarray[Any, Any]
+reveal_type(np.argsort(A, 0))  # E: ndarray[Any, dtype[{intp}]]
+reveal_type(np.argsort(B, 0))  # E: ndarray[Any, dtype[{intp}]]
 
 reveal_type(np.argmax(A))  # E: {intp}
 reveal_type(np.argmax(B))  # E: {intp}
@@ -76,38 +80,38 @@ reveal_type(np.argmin(B, axis=0))  # E: Any
 
 reveal_type(np.searchsorted(A[0], 0))  # E: {intp}
 reveal_type(np.searchsorted(B[0], 0))  # E: {intp}
-reveal_type(np.searchsorted(A[0], [0]))  # E: ndarray[Any, Any]
-reveal_type(np.searchsorted(B[0], [0]))  # E: ndarray[Any, Any]
+reveal_type(np.searchsorted(A[0], [0]))  # E: ndarray[Any, dtype[{intp}]]
+reveal_type(np.searchsorted(B[0], [0]))  # E: ndarray[Any, dtype[{intp}]]
 
-reveal_type(np.resize(a, (5, 5)))  # E: ndarray[Any, Any]
-reveal_type(np.resize(b, (5, 5)))  # E: ndarray[Any, Any]
-reveal_type(np.resize(c, (5, 5)))  # E: ndarray[Any, Any]
-reveal_type(np.resize(A, (5, 5)))  # E: ndarray[Any, Any]
-reveal_type(np.resize(B, (5, 5)))  # E: ndarray[Any, Any]
+reveal_type(np.resize(a, (5, 5)))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.resize(b, (5, 5)))  # E: ndarray[Any, dtype[{float32}]]
+reveal_type(np.resize(c, (5, 5)))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.resize(A, (5, 5)))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.resize(B, (5, 5)))  # E: ndarray[Any, dtype[{float32}]]
 
 reveal_type(np.squeeze(a))  # E: bool_
 reveal_type(np.squeeze(b))  # E: {float32}
-reveal_type(np.squeeze(c))  # E: ndarray[Any, Any]
-reveal_type(np.squeeze(A))  # E: ndarray[Any, Any]
-reveal_type(np.squeeze(B))  # E: ndarray[Any, Any]
+reveal_type(np.squeeze(c))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.squeeze(A))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.squeeze(B))  # E: ndarray[Any, dtype[{float32}]]
 
-reveal_type(np.diagonal(A))  # E: ndarray[Any, Any]
-reveal_type(np.diagonal(B))  # E: ndarray[Any, Any]
+reveal_type(np.diagonal(A))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.diagonal(B))  # E: ndarray[Any, dtype[{float32}]]
 
 reveal_type(np.trace(A))  # E: Any
 reveal_type(np.trace(B))  # E: Any
 
-reveal_type(np.ravel(a))  # E: ndarray[Any, Any]
-reveal_type(np.ravel(b))  # E: ndarray[Any, Any]
-reveal_type(np.ravel(c))  # E: ndarray[Any, Any]
-reveal_type(np.ravel(A))  # E: ndarray[Any, Any]
-reveal_type(np.ravel(B))  # E: ndarray[Any, Any]
+reveal_type(np.ravel(a))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.ravel(b))  # E: ndarray[Any, dtype[{float32}]]
+reveal_type(np.ravel(c))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.ravel(A))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.ravel(B))  # E: ndarray[Any, dtype[{float32}]]
 
-reveal_type(np.nonzero(a))  # E: tuple[ndarray[Any, Any]]
-reveal_type(np.nonzero(b))  # E: tuple[ndarray[Any, Any]]
-reveal_type(np.nonzero(c))  # E: tuple[ndarray[Any, Any]]
-reveal_type(np.nonzero(A))  # E: tuple[ndarray[Any, Any]]
-reveal_type(np.nonzero(B))  # E: tuple[ndarray[Any, Any]]
+reveal_type(np.nonzero(a))  # E: tuple[ndarray[Any, dtype[{intp}]]]
+reveal_type(np.nonzero(b))  # E: tuple[ndarray[Any, dtype[{intp}]]]
+reveal_type(np.nonzero(c))  # E: tuple[ndarray[Any, dtype[{intp}]]]
+reveal_type(np.nonzero(A))  # E: tuple[ndarray[Any, dtype[{intp}]]]
+reveal_type(np.nonzero(B))  # E: tuple[ndarray[Any, dtype[{intp}]]]
 
 reveal_type(np.shape(a))  # E: tuple[builtins.int]
 reveal_type(np.shape(b))  # E: tuple[builtins.int]
@@ -115,11 +119,11 @@ 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: ndarray[Any, Any]
-reveal_type(np.compress([True], b))  # E: ndarray[Any, Any]
-reveal_type(np.compress([True], c))  # E: ndarray[Any, Any]
-reveal_type(np.compress([True], A))  # E: ndarray[Any, Any]
-reveal_type(np.compress([True], B))  # E: ndarray[Any, Any]
+reveal_type(np.compress([True], a))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.compress([True], b))  # E: ndarray[Any, dtype[{float32}]]
+reveal_type(np.compress([True], c))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.compress([True], A))  # E: ndarray[Any, dtype[bool_]]
+reveal_type(np.compress([True], B))  # E: ndarray[Any, dtype[{float32}]]
 
 reveal_type(np.clip(a, 0, 1.0))  # E: Any
 reveal_type(np.clip(b, -1, 1))  # E: Any