Merge pull request #28750 from jorenham/typing/count_nonzero

charris · web-flow · commit b59d8ce8f583 · 2025-04-18T11:29:17.000-06:00
TYP: fix `count_nonzero` signature
diff --git a/numpy/_core/numeric.pyi b/numpy/_core/numeric.pyi
@@ -98,6 +98,7 @@ from numpy._typing import (
     _ArrayLikeComplex_co,
     _ArrayLikeTD64_co,
     _ArrayLikeObject_co,
+    _NestedSequence,
 )
 
 __all__ = [
@@ -413,21 +414,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
@@ -28,7 +28,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)
