hh.scalar_objects() strategy to refactor indexing test logic

Author: honno

array_api_tests/hypothesis_helpers.py

@@ -1,13 +1,23 @@
 import itertools
+import math
 from functools import reduce
-from math import sqrt
 from operator import mul
 from typing import Any, List, NamedTuple, Optional, Sequence, Tuple, Union
 
 from hypothesis import assume
-from hypothesis.strategies import (SearchStrategy, booleans, composite, floats,
-                                   integers, just, lists, none, one_of,
-                                   sampled_from, shared)
+from hypothesis.strategies import (
+    SearchStrategy,
+    booleans,
+    composite,
+    floats,
+    integers,
+    just,
+    lists,
+    none,
+    one_of,
+    sampled_from,
+    shared,
+)
 
 from . import _array_module as xp
 from . import dtype_helpers as dh
@@ -19,7 +29,7 @@
 from .algos import broadcast_shapes
 from .function_stubs import elementwise_functions
 from .pytest_helpers import nargs
-from .typing import Array, DataType, Shape
+from .typing import Array, DataType, Scalar, Shape
 
 # Set this to True to not fail tests just because a dtype isn't implemented.
 # If no compatible dtype is implemented for a given test, the test will fail
@@ -125,7 +135,7 @@ def mutually_promotable_dtypes(
 # Limit the total size of an array shape
 MAX_ARRAY_SIZE = 10000
 # Size to use for 2-dim arrays
-SQRT_MAX_ARRAY_SIZE = int(sqrt(MAX_ARRAY_SIZE))
+SQRT_MAX_ARRAY_SIZE = int(math.sqrt(MAX_ARRAY_SIZE))
 
 # np.prod and others have overflow and math.prod is Python 3.8+ only
 def prod(seq):
@@ -431,3 +441,11 @@ def axes(ndim: int) -> SearchStrategy[Optional[Union[int, Shape]]]:
         axes_strats.append(integers(-ndim, ndim - 1))
         axes_strats.append(xps.valid_tuple_axes(ndim))
     return one_of(axes_strats)
+
+
+def scalar_objects(dtype: DataType, shape: Shape) -> SearchStrategy[List[Scalar]]:
+    """Generates scalars or nested sequences which are valid for xp.asarray()"""
+    size = math.prod(shape)
+    return lists(xps.from_dtype(dtype), min_size=size, max_size=size).map(
+        lambda l: sh.reshape(l, shape)
+    )

array_api_tests/shape_helpers.py

@@ -1,9 +1,10 @@
+import math
 from itertools import product
 from typing import Iterator, List, Optional, Tuple, Union
 
-from .typing import Shape
+from .typing import Scalar, Shape
 
-__all__ = ["normalise_axis", "ndindex", "axis_ndindex", "axes_ndindex"]
+__all__ = ["normalise_axis", "ndindex", "axis_ndindex", "axes_ndindex", "reshape"]
 
 
 def normalise_axis(
@@ -57,3 +58,15 @@ def axes_ndindex(shape: Shape, axes: Tuple[int, ...]) -> Iterator[List[Shape]]:
             idx = tuple(idx)
             indices.append(idx)
         yield list(indices)
+
+
+def reshape(flat_seq: List[Scalar], shape: Shape) -> Union[Scalar, List[Scalar]]:
+    """Reshape a flat sequence"""
+    if len(shape) == 0:
+        assert len(flat_seq) == 1  # sanity check
+        return flat_seq[0]
+    elif len(shape) == 1:
+        return flat_seq
+    size = len(flat_seq)
+    n = math.prod(shape[1:])
+    return [reshape(flat_seq[i * n : (i + 1) * n], shape[1:]) for i in range(size // n)]

array_api_tests/test_array_object.py

@@ -1,6 +1,6 @@
 import math
 from itertools import product
-from typing import Sequence, Union, get_args
+from typing import get_args
 
 import pytest
 from hypothesis import assume, given, note
@@ -12,33 +12,13 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from .typing import DataType, Param, Scalar, ScalarType, Shape
-
-
-def reshape(
-    flat_seq: Sequence[Scalar], shape: Shape
-) -> Union[Scalar, Sequence[Scalar]]:
-    """Reshape a flat sequence"""
-    if len(shape) == 0:
-        assert len(flat_seq) == 1  # sanity check
-        return flat_seq[0]
-    elif len(shape) == 1:
-        return flat_seq
-    size = len(flat_seq)
-    n = math.prod(shape[1:])
-    return [reshape(flat_seq[i * n : (i + 1) * n], shape[1:]) for i in range(size // n)]
+from .typing import DataType, Param, Scalar, ScalarType
 
 
 @given(hh.shapes(min_side=1), st.data())  # TODO: test 0-sided arrays
 def test_getitem(shape, data):
-    size = math.prod(shape)
     dtype = data.draw(xps.scalar_dtypes(), label="dtype")
-    obj = data.draw(
-        st.lists(xps.from_dtype(dtype), min_size=size, max_size=size).map(
-            lambda l: reshape(l, shape)
-        ),
-        label="obj",
-    )
+    obj = data.draw(hh.scalar_objects(dtype, shape), label="obj")
     x = xp.asarray(obj, dtype=dtype)
     note(f"{x=}")
     key = data.draw(xps.indices(shape=shape), label="key")
@@ -71,21 +51,15 @@ def test_getitem(shape, data):
         for i in idx:
             val = val[i]
         out_obj.append(val)
-    out_obj = reshape(out_obj, out_shape)
+    out_obj = sh.reshape(out_obj, out_shape)
     expected = xp.asarray(out_obj, dtype=dtype)
     ph.assert_array("__getitem__", out, expected)
 
 
 @given(hh.shapes(min_side=1), st.data())  # TODO: test 0-sided arrays
 def test_setitem(shape, data):
-    size = math.prod(shape)
     dtype = data.draw(xps.scalar_dtypes(), label="dtype")
-    obj = data.draw(
-        st.lists(xps.from_dtype(dtype), min_size=size, max_size=size).map(
-            lambda l: reshape(l, shape)
-        ),
-        label="obj",
-    )
+    obj = data.draw(hh.scalar_objects(dtype, shape), label="obj")
     x = xp.asarray(obj, dtype=dtype)
     note(f"{x=}")
     # TODO: test setting non-0d arrays