Cover much more in test_setitem_mask

Author: honno

array_api_tests/pytest_helpers.py

@@ -204,18 +204,19 @@ def assert_scalar_equals(
     out: Scalar,
     expected: Scalar,
     /,
+    repr_name: str = "out",
     **kw,
 ):
-    out_repr = "out" if idx == () else f"out[{idx}]"
+    repr_name = repr_name if idx == () else f"{repr_name}[{idx}]"
     f_func = f"{func_name}({fmt_kw(kw)})"
     if type_ is bool or type_ is int:
-        msg = f"{out_repr}={out}, but should be {expected} [{f_func}]"
+        msg = f"{repr_name}={out}, but should be {expected} [{f_func}]"
         assert out == expected, msg
     elif math.isnan(expected):
-        msg = f"{out_repr}={out}, but should be {expected} [{f_func}]"
+        msg = f"{repr_name}={out}, but should be {expected} [{f_func}]"
         assert math.isnan(out), msg
     else:
-        msg = f"{out_repr}={out}, but should be roughly {expected} [{f_func}]"
+        msg = f"{repr_name}={out}, but should be roughly {expected} [{f_func}]"
         assert math.isclose(out, expected, rel_tol=0.25, abs_tol=1), msg
 
 

array_api_tests/test_array_object.py

@@ -156,16 +156,40 @@ def test_getitem_mask(shape, data):
                 )
 
 
-@given(hh.shapes(min_side=1), st.data())
+@given(hh.shapes(), st.data())
 def test_setitem_mask(shape, data):
     x = data.draw(xps.arrays(xps.scalar_dtypes(), shape=shape), label="x")
     key = data.draw(xps.arrays(dtype=xp.bool, shape=shape), label="key")
-    value = data.draw(xps.from_dtype(x.dtype), label="value")  # TODO: more values
+    value = data.draw(
+        xps.from_dtype(x.dtype) | xps.arrays(dtype=x.dtype, shape=()), label="value"
+    )
 
     res = xp.asarray(x, copy=True)
     res[key] = value
 
-    # TODO
+    ph.assert_dtype("__setitem__", x.dtype, res.dtype, repr_name="x.dtype")
+    ph.assert_shape("__setitem__", res.shape, x.shape, repr_name="x.dtype")
+
+    scalar_type = dh.get_scalar_type(x.dtype)
+    for idx in sh.ndindex(x.shape):
+        if key[idx]:
+            if isinstance(value, Scalar):
+                ph.assert_scalar_equals(
+                    "__setitem__",
+                    scalar_type,
+                    idx,
+                    scalar_type(res[idx]),
+                    value,
+                    repr_name="modified x",
+                )
+            else:
+                ph.assert_0d_equals(
+                    "__setitem__", "value", value, f"modified x[{idx}]", res[idx]
+                )
+        else:
+            ph.assert_0d_equals(
+                "__setitem__", f"old x[{idx}]", x[idx], f"modified x[{idx}]", res[idx]
+            )
 
 
 def make_param(method_name: str, dtype: DataType, stype: ScalarType) -> Param: