MAINT: simplifying annotations for np.core.from_numeric (#16556)

Simplified annotations for functions requiring >=1D ArrayLike
objects. Affects a set of <20 functions from
``np.core.fromnumeric``. Based on feedback from
https://github.com/numpy/numpy-stubs/pull/71.

3 files changed, 18 insertions, 47 deletions

diff --git a/.gitignore b/.gitignore
index a203abd5b..302968a14 100644
--- a/.gitignore
+++ b/.gitignore
@@ -28,6 +28,7 @@ GRTAGS
 GSYMS
 GTAGS
 .cache
+.mypy_cache/
 
 # Compiled source #
 ###################
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 5031893ed..2e9c28821 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -936,21 +936,18 @@ def reshape(a: ArrayLike, newshape: _ShapeLike, order: _Order = ...) -> ndarray:
 @overload
 def choose(
     a: _ScalarIntOrBool,
-    choices: Union[Sequence[ArrayLike], ndarray],
+    choices: ArrayLike,
     out: Optional[ndarray] = ...,
     mode: _Mode = ...,
 ) -> _ScalarIntOrBool: ...
 @overload
 def choose(
-    a: _IntOrBool,
-    choices: Union[Sequence[ArrayLike], ndarray],
-    out: Optional[ndarray] = ...,
-    mode: _Mode = ...,
+    a: _IntOrBool, choices: ArrayLike, out: Optional[ndarray] = ..., mode: _Mode = ...
 ) -> Union[integer, bool_]: ...
 @overload
 def choose(
     a: _ArrayLikeIntOrBool,
-    choices: Union[Sequence[ArrayLike], ndarray],
+    choices: ArrayLike,
     out: Optional[ndarray] = ...,
     mode: _Mode = ...,
 ) -> ndarray: ...
@@ -960,9 +957,7 @@ def repeat(
 def put(
     a: ndarray, ind: _ArrayLikeIntOrBool, v: ArrayLike, mode: _Mode = ...
 ) -> None: ...
-def swapaxes(
-    a: Union[Sequence[ArrayLike], ndarray], axis1: int, axis2: int
-) -> ndarray: ...
+def swapaxes(a: ArrayLike, axis1: int, axis2: int) -> ndarray: ...
 def transpose(
     a: ArrayLike, axes: Union[None, Sequence[int], ndarray] = ...
 ) -> ndarray: ...
@@ -998,54 +993,42 @@ def argpartition(
     order: Union[None, str, Sequence[str]] = ...,
 ) -> ndarray: ...
 def sort(
-    a: Union[Sequence[ArrayLike], ndarray],
+    a: ArrayLike,
     axis: Optional[int] = ...,
     kind: Optional[_SortKind] = ...,
     order: Union[None, str, Sequence[str]] = ...,
 ) -> ndarray: ...
 def argsort(
-    a: Union[Sequence[ArrayLike], ndarray],
+    a: ArrayLike,
     axis: Optional[int] = ...,
     kind: Optional[_SortKind] = ...,
     order: Union[None, str, Sequence[str]] = ...,
 ) -> ndarray: ...
 @overload
-def argmax(
-    a: Union[Sequence[ArrayLike], ndarray],
-    axis: None = ...,
-    out: Optional[ndarray] = ...,
-) -> integer: ...
+def argmax(a: ArrayLike, axis: None = ..., out: Optional[ndarray] = ...) -> integer: ...
 @overload
 def argmax(
-    a: Union[Sequence[ArrayLike], ndarray],
-    axis: int = ...,
-    out: Optional[ndarray] = ...,
+    a: ArrayLike, axis: int = ..., out: Optional[ndarray] = ...
 ) -> Union[integer, ndarray]: ...
 @overload
-def argmin(
-    a: Union[Sequence[ArrayLike], ndarray],
-    axis: None = ...,
-    out: Optional[ndarray] = ...,
-) -> integer: ...
+def argmin(a: ArrayLike, axis: None = ..., out: Optional[ndarray] = ...) -> integer: ...
 @overload
 def argmin(
-    a: Union[Sequence[ArrayLike], ndarray],
-    axis: int = ...,
-    out: Optional[ndarray] = ...,
+    a: ArrayLike, axis: int = ..., out: Optional[ndarray] = ...
 ) -> Union[integer, ndarray]: ...
 @overload
 def searchsorted(
-    a: Union[Sequence[ArrayLike], ndarray],
+    a: ArrayLike,
     v: _Scalar,
     side: _Side = ...,
-    sorter: Union[None, Sequence[_IntOrBool], ndarray] = ...,  # 1D int array
+    sorter: Optional[_ArrayLikeIntOrBool] = ...,  # 1D int array
 ) -> integer: ...
 @overload
 def searchsorted(
-    a: Union[Sequence[ArrayLike], ndarray],
+    a: ArrayLike,
     v: ArrayLike,
     side: _Side = ...,
-    sorter: Union[None, Sequence[_IntOrBool], ndarray] = ...,  # 1D int array
+    sorter: Optional[_ArrayLikeIntOrBool] = ...,  # 1D int array
 ) -> ndarray: ...
 def resize(a: ArrayLike, new_shape: _ShapeLike) -> ndarray: ...
 @overload
@@ -1053,13 +1036,10 @@ def squeeze(a: _ScalarGeneric, axis: Optional[_ShapeLike] = ...) -> _ScalarGener
 @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 = ...,
+    a: ArrayLike, offset: int = ..., axis1: int = ..., axis2: int = ...  # >= 2D array
 ) -> ndarray: ...
 def trace(
-    a: Union[Sequence[Sequence[ArrayLike]], ndarray],  # >= 2D array
+    a: ArrayLike,  # >= 2D array
     offset: int = ...,
     axis1: int = ...,
     axis2: int = ...,
@@ -1070,7 +1050,7 @@ 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
+    condition: ArrayLike,  # 1D bool array
     a: ArrayLike,
     axis: Optional[int] = ...,
     out: Optional[ndarray] = ...,
diff --git a/numpy/tests/typing/fail/fromnumeric.py b/numpy/tests/typing/fail/fromnumeric.py
index f158a1071..7455ce722 100644
--- a/numpy/tests/typing/fail/fromnumeric.py
+++ b/numpy/tests/typing/fail/fromnumeric.py
@@ -22,11 +22,9 @@ np.choose(A, mode="bob")  # E: No overload variant of "choose" matches argument
 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.swapaxes(a, 0, 0)  # E: Argument 1 to "swapaxes" has incompatible type
 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.transpose(a, axes=1)  # E: Argument "axes" to "transpose" has incompatible type
 np.transpose(A, axes=1.0)  # E: Argument "axes" to "transpose" has incompatible type
 
 np.partition(a, None)  # E: Argument 2 to "partition" has incompatible type
@@ -53,25 +51,20 @@ np.argpartition(
     A, 0, order=range(5)  # E: Argument "order" to "argpartition" has incompatible type
 )
 
-np.sort(a)  # E: Argument 1 to "sort" 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, order=range(5))  # E: Argument "order" to "sort" has incompatible type
 
-np.argsort(a)  # E: Argument 1 to "argsort" has incompatible type
 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"
 )
@@ -83,19 +76,16 @@ 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
 )