allow random int arrays

Author: flying-sheep

src/testing/fast_array_utils/_array_type.py

@@ -5,7 +5,7 @@
 
 import enum
 from dataclasses import KW_ONLY, dataclass, field
-from functools import cached_property
+from functools import cached_property, partial
 from importlib.metadata import version
 from typing import TYPE_CHECKING, Generic, TypeVar, cast
 
@@ -42,6 +42,9 @@ def __call__(
 
     _DTypeLikeFloat32 = np.dtype[np.float32] | type[np.float32]
     _DTypeLikeFloat64 = np.dtype[np.float64] | type[np.float64]
+    _DTypeLikeInt32 = np.dtype[np.int32] | type[np.int32]
+    _DTypeLikeIn64 = np.dtype[np.int64] | type[np.int64]
+    _DTypeLikeNum = _DTypeLikeFloat32 | _DTypeLikeFloat64 | _DTypeLikeInt32 | _DTypeLikeIn64
 else:
     Arr = TypeVar("Arr")
     Inner = TypeVar("Inner")
@@ -155,7 +158,7 @@ def random(
         self,
         shape: tuple[int, int],
         *,
-        dtype: _DTypeLikeFloat32 | _DTypeLikeFloat64 | None = None,
+        dtype: _DTypeLikeNum | None = None,
         gen: np.random.Generator | None = None,
         # sparse only
         density: float | np.floating[Any] = 0.01,
@@ -165,7 +168,7 @@ def random(
 
         match self.mod, self.name, self.inner:
             case "numpy", "ndarray", None:
-                return cast("Arr", gen.random(shape, dtype=dtype or np.float64))
+                return cast("Arr", random_array(shape, dtype=dtype, rng=gen))
             case "scipy.sparse", (
                 "csr_array" | "csc_array" | "csr_matrix" | "csc_matrix"
             ) as cls_name, None:
@@ -179,7 +182,7 @@ def random(
                     ),
                 )
             case "cupy", "ndarray", None:
-                return self(gen.random(shape, dtype=dtype or np.float64))
+                return self(random_array(shape, dtype=dtype, rng=gen))
             case "cupyx.scipy.sparse", ("csr_matrix" | "csc_matrix") as cls_name, None:
                 import cupy as cu
 
@@ -363,12 +366,28 @@ def _to_cupy_sparse(
         return self.cls(x)  # type: ignore[call-arg,arg-type, return-value]
 
 
+def random_array(
+    shape: tuple[int, int],
+    *,
+    dtype: _DTypeLikeNum | None = None,
+    rng: np.random.Generator | None = None,
+) -> Array:
+    """Create a random array."""
+    rng = np.random.default_rng(rng)
+    f = (
+        partial(rng.integers, 0, 10_000)
+        if dtype is not None and np.dtype(dtype).kind in "iu"
+        else rng.random
+    )
+    return f(shape, dtype=dtype)  # type: ignore[arg-type]
+
+
 def random_mat(
     shape: tuple[int, int],
     *,
     density: float | np.floating[Any] = 0.01,
     format: Literal["csr", "csc"] = "csr",  # noqa: A002
-    dtype: _DTypeLikeFloat32 | _DTypeLikeFloat64 | None = None,
+    dtype: _DTypeLikeNum | None = None,
     container: Literal["array", "matrix"] = "array",
     rng: np.random.Generator | None = None,
 ) -> types.CSBase:

tests/test_sparse.py

@@ -19,7 +19,7 @@
 
     from fast_array_utils.types import CSBase
     from testing.fast_array_utils import ArrayType
-    from testing.fast_array_utils._array_type import _DTypeLikeFloat32, _DTypeLikeFloat64
+    from testing.fast_array_utils._array_type import _DTypeLikeNum
 
 
 pytestmark = [pytest.mark.skipif(not find_spec("scipy"), reason="scipy not installed")]
@@ -48,9 +48,7 @@ def dtype(request: pytest.FixtureRequest) -> type[np.float32 | np.float64]:
 @pytest.mark.array_type(select=Flags.Sparse, skip=Flags.Dask | Flags.Disk | Flags.Gpu)
 @pytest.mark.parametrize("order", ["C", "F"])
 def test_to_dense(
-    array_type: ArrayType[CSBase, None],
-    order: Literal["C", "F"],
-    dtype: _DTypeLikeFloat32 | _DTypeLikeFloat64,
+    array_type: ArrayType[CSBase, None], order: Literal["C", "F"], dtype: _DTypeLikeNum
 ) -> None:
     mat = array_type.random((10, 10), density=0.1, dtype=dtype)
     with WARNS_NUMBA if not find_spec("numba") else nullcontext():
@@ -68,7 +66,7 @@ def test_to_dense_benchmark(
     benchmark: BenchmarkFixture,
     array_type: ArrayType[CSBase, None],
     order: Literal["C", "F"],
-    dtype: _DTypeLikeFloat32 | _DTypeLikeFloat64,
+    dtype: _DTypeLikeNum,
 ) -> None:
     mat = array_type.random((10_000, 10_000), dtype=dtype)
     to_dense(mat, order=order)  # warmup: numba compile

tests/test_stats.py

@@ -199,7 +199,7 @@ def test_dask_constant_blocks(
 @pytest.mark.benchmark
 @pytest.mark.array_type(skip=Flags.Matrix | Flags.Dask | Flags.Disk | Flags.Gpu)
 @pytest.mark.parametrize("func", [stats.sum, stats.mean, stats.mean_var, stats.is_constant])
-@pytest.mark.parametrize("dtype", [np.float32, np.float64])  # random only supports float
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.int32])
 def test_stats_benchmark(
     benchmark: BenchmarkFixture,
     func: BenchFun,