🐛 sparse: fix csc and csr return type of count_nonzero with axis given

Author: jorenham

scipy-stubs/sparse/_csc.pyi

@@ -1,7 +1,9 @@
 from typing import Any, Generic, Literal, overload
 from typing_extensions import TypeIs, TypeVar, override
 
+import numpy as np
 import optype as op
+import optype.numpy as onp
 from ._base import sparray
 from ._compressed import _cs_matrix
 from ._matrix import spmatrix
@@ -24,6 +26,12 @@ class _csc_base(_cs_matrix[_SCT, tuple[int, int]], Generic[_SCT]):
     @override
     def shape(self, /) -> tuple[int, int]: ...
 
+    #
+    @overload  # type: ignore[explicit-override]
+    def count_nonzero(self, /, axis: None = None) -> int: ...
+    @overload
+    def count_nonzero(self, /, axis: op.CanIndex) -> onp.Array1D[np.intp]: ...
+
 class csc_array(_csc_base[_SCT], sparray, Generic[_SCT]): ...
 
 class csc_matrix(_csc_base[_SCT], spmatrix[_SCT], Generic[_SCT]):  # type: ignore[misc]

scipy-stubs/sparse/_csr.pyi

@@ -1,7 +1,9 @@
 from typing import Any, Generic, Literal, overload
 from typing_extensions import TypeIs, TypeVar, override
 
+import numpy as np
 import optype as op
+import optype.numpy as onp
 from ._base import sparray
 from ._compressed import _cs_matrix
 from ._matrix import spmatrix
@@ -22,6 +24,16 @@ class _csr_base(_cs_matrix[_SCT, _ShapeT_co], Generic[_SCT, _ShapeT_co]):
     @override
     def format(self, /) -> Literal["csr"]: ...
 
+    #
+    @overload  # type: ignore[explicit-override]
+    def count_nonzero(self, /, axis: None = None) -> int: ...
+    @overload
+    def count_nonzero(self: _csr_base[Any, tuple[int]], /, axis: op.CanIndex) -> int: ...
+    @overload
+    def count_nonzero(self: _csr_base[Any, tuple[int, int]], /, axis: op.CanIndex) -> onp.Array1D[np.intp]: ...
+    @overload
+    def count_nonzero(self: csr_array, /, axis: op.CanIndex) -> int | onp.Array1D[np.intp]: ...  # type: ignore[misc]
+
 class csr_array(_csr_base[_SCT, _ShapeT_co], sparray, Generic[_SCT, _ShapeT_co]): ...
 
 class csr_matrix(_csr_base[_SCT, tuple[int, int]], spmatrix[_SCT], Generic[_SCT]):  # type: ignore[misc]