Test non-0d-resulting keys in test_setitem

Also `ph.assert_array()` -> `ph.assert_array_elements()`

Author: honno

array_api_tests/meta/test_pytest_helpers.py

@@ -13,10 +13,10 @@ def test_assert_dtype():
     ph.assert_dtype("single_bool_func", [xp.uint8], xp.bool, xp.bool)
 
 
-def test_assert_array():
-    ph.assert_array("int zeros", xp.asarray(0), xp.asarray(0))
-    ph.assert_array("pos zeros", xp.asarray(0.0), xp.asarray(0.0))
+def test_assert_array_elements():
+    ph.assert_array_elements("int zeros", xp.asarray(0), xp.asarray(0))
+    ph.assert_array_elements("pos zeros", xp.asarray(0.0), xp.asarray(0.0))
     with raises(AssertionError):
-        ph.assert_array("mixed sign zeros", xp.asarray(0.0), xp.asarray(-0.0))
+        ph.assert_array_elements("mixed sign zeros", xp.asarray(0.0), xp.asarray(-0.0))
     with raises(AssertionError):
-        ph.assert_array("mixed sign zeros", xp.asarray(-0.0), xp.asarray(0.0))
+        ph.assert_array_elements("mixed sign zeros", xp.asarray(-0.0), xp.asarray(0.0))

array_api_tests/pytest_helpers.py

@@ -25,7 +25,7 @@
     "assert_keepdimable_shape",
     "assert_0d_equals",
     "assert_fill",
-    "assert_array",
+    "assert_array_elements",
 ]
 
 
@@ -374,28 +374,30 @@ def assert_fill(
         assert xp.all(xp.equal(out, xp.asarray(fill_value, dtype=dtype))), msg
 
 
-def assert_array(func_name: str, out: Array, expected: Array, /, **kw):
+def assert_array_elements(
+    func_name: str, out: Array, expected: Array, /, *, out_repr: str = "out", **kw
+):
     """
-    Assert array is (strictly) as expected, e.g.
+    Assert array elements are (strictly) as expected, e.g.
 
         >>> x = xp.arange(5)
         >>> out = xp.asarray(x)
-        >>> assert_array('asarray', out, x)
+        >>> assert_array_elements('asarray', out, x)
 
         is equivalent to
 
         >>> assert xp.all(out == x)
 
     """
-    assert_dtype(func_name, out.dtype, expected.dtype)
-    assert_shape(func_name, out.shape, expected.shape, **kw)
+    dh.result_type(out.dtype, expected.dtype)  # sanity check
+    assert_shape(func_name, out.shape, expected.shape, **kw)  # sanity check
     f_func = f"[{func_name}({fmt_kw(kw)})]"
     if dh.is_float_dtype(out.dtype):
         for idx in sh.ndindex(out.shape):
             at_out = out[idx]
             at_expected = expected[idx]
             msg = (
-                f"{sh.fmt_idx('out', idx)}={at_out}, should be {at_expected} "
+                f"{sh.fmt_idx(out_repr, idx)}={at_out}, should be {at_expected} "
                 f"{f_func}"
             )
             if xp.isnan(at_expected):
@@ -411,6 +413,6 @@ def assert_array(func_name: str, out: Array, expected: Array, /, **kw):
             else:
                 assert at_out == at_expected, msg
     else:
-        assert xp.all(out == expected), (
-            f"out not as expected {f_func}\n" f"{out=}\n{expected=}"
-        )
+        assert xp.all(
+            out == expected
+        ), f"{out_repr} not as expected {f_func}\n{out_repr}={out!r}\n{expected=}"

array_api_tests/test_array_object.py

@@ -25,12 +25,11 @@ def scalar_objects(dtype: DataType, shape: Shape) -> st.SearchStrategy[List[Scal
     )
 
 
-@given(hh.shapes(), st.data())
-def test_getitem(shape, data):
-    dtype = data.draw(xps.scalar_dtypes(), label="dtype")
+@given(shape=hh.shapes(), dtype=xps.scalar_dtypes(), data=st.data())
+def test_getitem(shape, dtype, data):
     zero_sided = any(side == 0 for side in shape)
     if zero_sided:
-        x = xp.ones(shape, dtype=dtype)
+        x = xp.zeros(shape, dtype=dtype)
     else:
         obj = data.draw(scalar_objects(dtype, shape), label="obj")
         x = xp.asarray(obj, dtype=dtype)
@@ -76,45 +75,62 @@ def test_getitem(shape, data):
             out_obj.append(val)
         out_obj = sh.reshape(out_obj, out_shape)
         expected = xp.asarray(out_obj, dtype=dtype)
-        ph.assert_array("__getitem__", out, expected)
+        ph.assert_array_elements("__getitem__", out, expected)
 
 
-@given(hh.shapes(min_side=1), st.data())  # TODO: test 0-sided arrays
-def test_setitem(shape, data):
-    dtype = data.draw(xps.scalar_dtypes(), label="dtype")
-    obj = data.draw(scalar_objects(dtype, shape), label="obj")
-    x = xp.asarray(obj, dtype=dtype)
+@given(shape=hh.shapes(min_side=1), dtype=xps.scalar_dtypes(), data=st.data())
+def test_setitem(shape, dtype, data):
+    zero_sided = any(side == 0 for side in shape)
+    if zero_sided:
+        x = xp.zeros(shape, dtype=dtype)
+    else:
+        obj = data.draw(scalar_objects(dtype, shape), label="obj")
+        x = xp.asarray(obj, dtype=dtype)
     note(f"{x=}")
-    # TODO: test setting non-0d arrays
-    key = data.draw(xps.indices(shape=shape, max_dims=0), label="key")
-    value = data.draw(
-        xps.from_dtype(dtype) | xps.arrays(dtype=dtype, shape=()), label="value"
-    )
+    key = data.draw(xps.indices(shape=shape), label="key")
+    _key = tuple(key) if isinstance(key, tuple) else (key,)
+    if Ellipsis in _key:
+        nonexpanding_key = tuple(i for i in _key if i is not None)
+        start_a = nonexpanding_key.index(Ellipsis)
+        stop_a = start_a + (len(shape) - (len(nonexpanding_key) - 1))
+        slices = tuple(slice(None) for _ in range(start_a, stop_a))
+        start_pos = _key.index(Ellipsis)
+        _key = _key[:start_pos] + slices + _key[start_pos + 1 :]
+    out_shape = []
+    for a, i in enumerate(_key):
+        if isinstance(i, slice):
+            side = shape[a]
+            indices = range(side)[i]
+            out_shape.append(len(indices))
+    out_shape = tuple(out_shape)
+    value_strat = xps.arrays(dtype=dtype, shape=out_shape)
+    if out_shape == ():
+        # We can pass scalars if we're only indexing one element
+        value_strat |= xps.from_dtype(dtype)
+    value = data.draw(value_strat, label="value")
 
     res = xp.asarray(x, copy=True)
     res[key] = value
 
     ph.assert_dtype("__setitem__", x.dtype, res.dtype, repr_name="x.dtype")
     ph.assert_shape("__setitem__", res.shape, x.shape, repr_name="x.shape")
+    f_res = f"res[{sh.fmt_idx('x', key)}]"
     if isinstance(value, get_args(Scalar)):
-        msg = f"x[{key}]={res[key]!r}, but should be {value=} [__setitem__()]"
+        msg = f"{f_res}={res[key]!r}, but should be {value=} [__setitem__()]"
         if math.isnan(value):
             assert xp.isnan(res[key]), msg
         else:
             assert res[key] == value, msg
     else:
-        ph.assert_0d_equals(
-            "__setitem__", "value", value, f"modified x[{key}]", res[key]
-        )
-    _key = key if isinstance(key, tuple) else (key,)
-    assume(all(isinstance(i, int) for i in _key))  # TODO: normalise slices and ellipsis
-    _key = tuple(i if i >= 0 else s + i for i, s in zip(_key, x.shape))
-    unaffected_indices = list(sh.ndindex(res.shape))
-    unaffected_indices.remove(_key)
-    for idx in unaffected_indices:
-        ph.assert_0d_equals(
-            "__setitem__", f"old x[{idx}]", x[idx], f"modified x[{idx}]", res[idx]
-        )
+        ph.assert_array_elements("__setitem__", res[key], value, out_repr=f_res)
+    if all(isinstance(i, int) for i in _key):  # TODO: normalise slices and ellipsis
+        _key = tuple(i if i >= 0 else s + i for i, s in zip(_key, x.shape))
+        unaffected_indices = list(sh.ndindex(res.shape))
+        unaffected_indices.remove(_key)
+        for idx in unaffected_indices:
+            ph.assert_0d_equals(
+                "__setitem__", f"old x[{idx}]", x[idx], f"modified x[{idx}]", res[idx]
+            )
 
 
 @pytest.mark.data_dependent_shapes

array_api_tests/test_creation_functions.py

@@ -180,7 +180,7 @@ def test_arange(dtype, data):
     if dh.is_int_dtype(_dtype):
         elements = list(r)
         assume(out.size == len(elements))
-        ph.assert_array("arange", out, xp.asarray(elements, dtype=_dtype))
+        ph.assert_array_elements("arange", out, xp.asarray(elements, dtype=_dtype))
     else:
         assume(out.size == size)
         if out.size > 0:
@@ -262,7 +262,7 @@ def test_asarray_arrays(x, data):
         ph.assert_kw_dtype("asarray", dtype, out.dtype)
     ph.assert_shape("asarray", out.shape, x.shape)
     if dtype is None or dtype == x.dtype:
-        ph.assert_array("asarray", out, x, **kw)
+        ph.assert_array_elements("asarray", out, x, **kw)
     else:
         pass  # TODO
     copy = kw.get("copy", None)
@@ -452,7 +452,7 @@ def test_linspace(num, dtype, endpoint, data):
         # the first num elements when endpoint=False
         expected = xp.linspace(start, stop, num + 1, dtype=dtype, endpoint=True)
         expected = expected[:-1]
-        ph.assert_array("linspace", out, expected)
+        ph.assert_array_elements("linspace", out, expected)
 
 
 @given(dtype=xps.numeric_dtypes(), data=st.data())

array_api_tests/test_operators_and_elementwise_functions.py

@@ -1124,7 +1124,7 @@ def test_positive(ctx, data):
 
     ph.assert_dtype(ctx.func_name, x.dtype, out.dtype)
     ph.assert_shape(ctx.func_name, out.shape, x.shape)
-    ph.assert_array(ctx.func_name, out, x)
+    ph.assert_array_elements(ctx.func_name, out, x)
 
 
 @pytest.mark.parametrize("ctx", make_binary_params("pow", dh.numeric_dtypes))