Merge pull request numpy#27208 from jorenham/typing/nonzero-scalar

TYP: Disallow scalars and 0d-arrays in ``numpy.nonzero``

Author: seberg

numpy/_core/fromnumeric.pyi

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Any, overload, TypeVar, Literal, SupportsIndex
+from typing import Any, NoReturn, overload, TypeVar, Literal, SupportsIndex
 
 import numpy as np
 from numpy import (
@@ -368,7 +368,10 @@ def ravel(a: _ArrayLike[_SCT], order: _OrderKACF = ...) -> NDArray[_SCT]: ...
 @overload
 def ravel(a: ArrayLike, order: _OrderKACF = ...) -> NDArray[Any]: ...
 
-def nonzero(a: ArrayLike) -> tuple[NDArray[intp], ...]: ...
+@overload
+def nonzero(a: np.generic | np.ndarray[tuple[()], Any]) -> NoReturn: ...
+@overload
+def nonzero(a: _ArrayLike[Any]) -> tuple[NDArray[intp], ...]: ...
 
 def shape(a: ArrayLike) -> _Shape: ...
 

numpy/typing/tests/data/fail/fromnumeric.pyi

@@ -89,6 +89,8 @@ np.trace(A, axis2=[])  # E: No overload variant
 
 np.ravel(a, order="bob")  # E: No overload variant
 
+np.nonzero(0)  # E: No overload variant
+
 np.compress(  # E: No overload variant
     [True], A, axis=1.0
 )

numpy/typing/tests/data/reveal/fromnumeric.pyi

@@ -1,7 +1,7 @@
 """Tests for :mod:`_core.fromnumeric`."""
 
 import sys
-from typing import Any
+from typing import Any, NoReturn
 
 import numpy as np
 import numpy.typing as npt
@@ -22,6 +22,9 @@ AR_i8: npt.NDArray[np.int64]
 AR_O: npt.NDArray[np.object_]
 AR_subclass: NDArraySubclass
 AR_m: npt.NDArray[np.timedelta64]
+AR_0d: np.ndarray[tuple[()], np.dtype[Any]]
+AR_1d: np.ndarray[tuple[int], np.dtype[Any]]
+AR_nd: np.ndarray[tuple[int, ...], np.dtype[Any]]
 
 b: np.bool
 f4: np.float32
@@ -128,11 +131,13 @@ assert_type(np.ravel(f), npt.NDArray[Any])
 assert_type(np.ravel(AR_b), npt.NDArray[np.bool])
 assert_type(np.ravel(AR_f4), npt.NDArray[np.float32])
 
-assert_type(np.nonzero(b), tuple[npt.NDArray[np.intp], ...])
-assert_type(np.nonzero(f4), tuple[npt.NDArray[np.intp], ...])
-assert_type(np.nonzero(f), tuple[npt.NDArray[np.intp], ...])
+assert_type(np.nonzero(b), NoReturn)
+assert_type(np.nonzero(f4), NoReturn)
 assert_type(np.nonzero(AR_b), tuple[npt.NDArray[np.intp], ...])
 assert_type(np.nonzero(AR_f4), tuple[npt.NDArray[np.intp], ...])
+assert_type(np.nonzero(AR_0d), NoReturn)
+assert_type(np.nonzero(AR_1d), tuple[npt.NDArray[np.intp], ...])
+assert_type(np.nonzero(AR_nd), tuple[npt.NDArray[np.intp], ...])
 
 assert_type(np.shape(b), tuple[int, ...])
 assert_type(np.shape(f4), tuple[int, ...])