fix types

Author: flying-sheep

src/fast_array_utils/stats/_typing.py

@@ -34,7 +34,7 @@ def __call__(
         *,
         axis: Literal[0, 1] | None = None,
         keep_cupy_as_array: bool = False,
-    ) -> types.DaskArray | types.HasArrayNamespace: ...
+    ) -> NDArray[Any] | types.CupyArray | np.number[Any] | types.DaskArray | types.HasArrayNamespace: ...
 
 
 class StatFunDtype(Protocol):

src/fast_array_utils/types.py

@@ -7,6 +7,10 @@
 from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 
+if TYPE_CHECKING:
+    from types import ModuleType
+
+
 __all__ = [
     "COOBase",
     "CSArray",
@@ -22,6 +26,7 @@
     "DaskArray",
     "H5Dataset",
     "H5Group",
+    "HasArrayNamespace",
     "ZarrArray",
     "ZarrGroup",
     "coo_array",
@@ -37,8 +42,6 @@
 
 # scipy sparse
 if TYPE_CHECKING:
-    from types import ModuleType
-
     from scipy.sparse import coo_array, coo_matrix, csc_array, csc_matrix, csr_array, csr_matrix, sparray, spmatrix
 else:
     try:  # cs?_array isn’t available in older scipy versions
@@ -124,7 +127,17 @@
 class HasArrayNamespace(Protocol):
     """An array object compatible with the Python array API standard."""
 
-    ndim: int
-    shape: tuple[int, ...]
+    @property
+    def ndim(self) -> int:
+        """The number of dimensions of the array."""
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """The shape of the array."""
+
+    @property
+    def dtype(self) -> object:
+        """The data type of the array."""
 
-    def __array_namespace__(self, /, *, api_version: str | None = None) -> ModuleType: ...
+    def __array_namespace__(self, /, *, api_version: str | None = None) -> ModuleType:
+        """Get Array API namespace."""

tests/test_jax.py

@@ -14,6 +14,7 @@
 if TYPE_CHECKING:
     from typing import Literal
 
+
 pytestmark = pytest.mark.skipif(not find_spec("jax"), reason="jax not installed")
 
 if find_spec("jax"):

tests/test_stats.py

@@ -103,9 +103,11 @@ def np_arr(dtype_in: type[DTypeIn], ndim: Literal[1, 2]) -> NDArray[DTypeIn]:
     return np_arr
 
 
-def to_np_dense_checked(
-    stat: NDArray[DTypeOut] | np.number[Any] | types.DaskArray, axis: Literal[0, 1] | None, arr: CpuArray | GpuArray | DiskArray | types.DaskArray
-) -> NDArray[DTypeOut] | np.number[Any]:
+def to_np_dense_checked[DT: DTypeOut](
+    stat: NDArray[DT] | np.number[Any] | types.DaskArray | types.HasArrayNamespace,
+    axis: Literal[0, 1] | None,
+    arr: CpuArray | GpuArray | DiskArray | types.COOBase | types.DaskArray | types.HasArrayNamespace,
+) -> NDArray[DT] | np.number[Any]:
     match axis, arr:
         case _, types.DaskArray():
             assert isinstance(stat, types.DaskArray), type(stat)
@@ -208,7 +210,7 @@ def test_min_max(array_type: ArrayType[CpuArray | GpuArray | DiskArray | types.D
     np_arr = rng.random((100, 100))
     arr = array_type(np_arr)
 
-    result = to_np_dense_checked(func(arr, axis=axis), axis, arr)
+    result = to_np_dense_checked(func(arr, axis=axis), axis, arr)  # type: ignore[arg-type]
 
     expected = (np.min if func is stats.min else np.max)(np_arr, axis=axis)
     np.testing.assert_array_equal(result, expected)
@@ -229,7 +231,7 @@ def test_dask_shapes(array_type: ArrayType[types.DaskArray], axis: Literal[0, 1]
     np_arr = np.array(data, dtype=np.float32)
     arr = array_type(np_arr)
     assert 1 in arr.chunksize, "This test is supposed to test 1×n and n×1 chunk sizes"
-    stat = cast("NDArray[Any] | types.CupyArray", func(arr, axis=axis).compute())
+    stat = cast("NDArray[Any] | types.CupyArray", func(arr, axis=axis).compute())  # type: ignore[union-attr]
     if isinstance(stat, types.CupyArray):
         stat = stat.get()
     np_func = getattr(np, func.__name__)
@@ -321,6 +323,8 @@ def test_mean_var_pbmc_dask(array_type: ArrayType[types.DaskArray], pbmc64k_redu
     arr = array_type(mat)
 
     mean_mat, var_mat = stats.mean_var(mat, axis=0, correction=1)
+    mean_arr: NDArray[Any] | np.number  # actually just NDArray, and mypy should be able to infer.
+    var_arr: NDArray[Any] | np.number
     mean_arr, var_arr = (to_np_dense_checked(a, 0, arr) for a in stats.mean_var(arr, axis=0, correction=1))
 
     rtol = 1.0e-5 if array_type.flags & Flags.Gpu else 1.0e-7

typings/cupyx/scipy/sparse/_coo.pyi

@@ -8,4 +8,4 @@ from ._base import spmatrix
 
 class coo_matrix(spmatrix):
     format: Literal["coo"] = "coo"
-    def get(self, stream: cupy.cuda.Stream | None = None) -> sps.spmatrix: ...
+    def get(self, stream: cupy.cuda.Stream | None = None) -> sps.coo_matrix: ...

typings/cupyx/scipy/sparse/_csr.pyi

@@ -8,4 +8,4 @@ from ._compressed import _compressed_sparse_matrix
 
 class csr_matrix(_compressed_sparse_matrix):
     format: Literal["csr"] = "csr"
-    def get(self, stream: cupy.cuda.Stream | None = None) -> sps.csc_matrix: ...
+    def get(self, stream: cupy.cuda.Stream | None = None) -> sps.csr_matrix: ...