Merge pull request #28774 from charris/backport-28750

charris · web-flow · commit b3e772cc081f · 2025-04-19T07:17:56.000-06:00
TYP: fix `count_nonzero` signature
diff --git a/numpy/_core/numeric.pyi b/numpy/_core/numeric.pyi
@@ -110,6 +110,7 @@ from numpy._typing import (
     _ArrayLikeTD64_co,
     _ArrayLikeObject_co,
     _ArrayLikeUnknown,
+    _NestedSequence,
 )
 
 __all__ = [
@@ -456,21 +457,19 @@ def full_like(
     device: None | L["cpu"] = ...,
 ) -> NDArray[Any]: ...
 
+#
 @overload
-def count_nonzero(
-    a: ArrayLike,
-    axis: None = ...,
-    *,
-    keepdims: L[False] = ...,
-) -> int: ...
+def count_nonzero(a: ArrayLike, axis: None = None, *, keepdims: L[False] = False) -> int: ...
+@overload
+def count_nonzero(a: _ScalarLike_co, axis: _ShapeLike | None = None, *, keepdims: L[True]) -> np.intp: ...
 @overload
 def count_nonzero(
-    a: ArrayLike,
-    axis: _ShapeLike = ...,
-    *,
-    keepdims: bool = ...,
-) -> Any: ...  # TODO: np.intp or ndarray[np.intp]
+    a: NDArray[Any] | _NestedSequence[ArrayLike], axis: _ShapeLike | None = None, *, keepdims: L[True]
+) -> NDArray[np.intp]: ...
+@overload
+def count_nonzero(a: ArrayLike, axis: _ShapeLike | None = None, *, keepdims: bool = False) -> Any: ...
 
+#
 def isfortran(a: NDArray[Any] | generic) -> bool: ...
 
 def argwhere(a: ArrayLike) -> NDArray[intp]: ...
diff --git a/numpy/typing/tests/data/reveal/numeric.pyi b/numpy/typing/tests/data/reveal/numeric.pyi
@@ -31,7 +31,7 @@ C: SubClass
 assert_type(np.count_nonzero(i8), int)
 assert_type(np.count_nonzero(AR_i8), int)
 assert_type(np.count_nonzero(B), int)
-assert_type(np.count_nonzero(AR_i8, keepdims=True), Any)
+assert_type(np.count_nonzero(AR_i8, keepdims=True), npt.NDArray[np.intp])
 assert_type(np.count_nonzero(AR_i8, axis=0), Any)
 
 assert_type(np.isfortran(i8), bool)
