ENH: add "repro_snippets" to test_utility_functions.py

Author: ev-br

array_api_tests/test_utility_functions.py

@@ -18,23 +18,28 @@ def test_all(x, data):
     kw = data.draw(hh.kwargs(axis=hh.axes(x.ndim), keepdims=st.booleans()), label="kw")
     keepdims = kw.get("keepdims", False)
 
-    out = xp.all(x, **kw)
-
-    ph.assert_dtype("all", in_dtype=x.dtype, out_dtype=out.dtype, expected=xp.bool)
-    _axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
-    ph.assert_keepdimable_shape(
-        "all", in_shape=x.shape, out_shape=out.shape, axes=_axes, keepdims=keepdims, kw=kw
-    )
-    scalar_type = dh.get_scalar_type(x.dtype)
-    for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
-        result = bool(out[out_idx])
-        elements = []
-        for idx in indices:
-            s = scalar_type(x[idx])
-            elements.append(s)
-        expected = all(elements)
-        ph.assert_scalar_equals("all", type_=scalar_type, idx=out_idx,
-                                out=result, expected=expected, kw=kw)
+    repro_snippet = ph.format_snippet(f"xp.all({x!r}, **kw) with {kw = }")
+    try:
+        out = xp.all(x, **kw)
+
+        ph.assert_dtype("all", in_dtype=x.dtype, out_dtype=out.dtype, expected=xp.bool)
+        _axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
+        ph.assert_keepdimable_shape(
+            "all", in_shape=x.shape, out_shape=out.shape, axes=_axes, keepdims=keepdims, kw=kw
+        )
+        scalar_type = dh.get_scalar_type(x.dtype)
+        for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
+            result = bool(out[out_idx])
+            elements = []
+            for idx in indices:
+                s = scalar_type(x[idx])
+                elements.append(s)
+            expected = all(elements)
+            ph.assert_scalar_equals("all", type_=scalar_type, idx=out_idx,
+                                    out=result, expected=expected, kw=kw)
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise
 
 
 @pytest.mark.unvectorized
@@ -46,23 +51,28 @@ def test_any(x, data):
     kw = data.draw(hh.kwargs(axis=hh.axes(x.ndim), keepdims=st.booleans()), label="kw")
     keepdims = kw.get("keepdims", False)
 
-    out = xp.any(x, **kw)
-
-    ph.assert_dtype("any", in_dtype=x.dtype, out_dtype=out.dtype, expected=xp.bool)
-    _axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
-    ph.assert_keepdimable_shape(
-        "any", in_shape=x.shape, out_shape=out.shape, axes=_axes, keepdims=keepdims, kw=kw,
-    )
-    scalar_type = dh.get_scalar_type(x.dtype)
-    for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
-        result = bool(out[out_idx])
-        elements = []
-        for idx in indices:
-            s = scalar_type(x[idx])
-            elements.append(s)
-        expected = any(elements)
-        ph.assert_scalar_equals("any", type_=scalar_type, idx=out_idx,
-                                out=result, expected=expected, kw=kw)
+    repro_snippet = ph.format_snippet(f"xp.any({x!r}, **kw) with {kw = }")
+    try:
+        out = xp.any(x, **kw)
+
+        ph.assert_dtype("any", in_dtype=x.dtype, out_dtype=out.dtype, expected=xp.bool)
+        _axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
+        ph.assert_keepdimable_shape(
+            "any", in_shape=x.shape, out_shape=out.shape, axes=_axes, keepdims=keepdims, kw=kw,
+        )
+        scalar_type = dh.get_scalar_type(x.dtype)
+        for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
+            result = bool(out[out_idx])
+            elements = []
+            for idx in indices:
+                s = scalar_type(x[idx])
+                elements.append(s)
+            expected = any(elements)
+            ph.assert_scalar_equals("any", type_=scalar_type, idx=out_idx,
+                                    out=result, expected=expected, kw=kw)
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise
 
 
 @pytest.mark.unvectorized
@@ -85,19 +95,24 @@ def test_diff(x, data):
 
     n = data.draw(st.integers(1, min(x.shape[n_axis], 3)))
 
-    out = xp.diff(x, **axis_kw, n=n)
+    repro_snippet = ph.format_snippet(f"xp.diff({x!r}, **axis_kw, n={n!r}) with {axis_kw = }")
+    try:
+        out = xp.diff(x, **axis_kw, n=n)
 
-    expected_shape = list(x.shape)
-    expected_shape[n_axis] -= n
+        expected_shape = list(x.shape)
+        expected_shape[n_axis] -= n
 
-    assert out.shape == tuple(expected_shape)
+        assert out.shape == tuple(expected_shape)
 
-    # value test
-    if n == 1:
-        for idx in sh.ndindex(out.shape):
-            l = list(idx)
-            l[n_axis] += 1
-            assert out[idx] == x[tuple(l)] - x[idx], f"diff failed with {idx = }"
+        # value test
+        if n == 1:
+            for idx in sh.ndindex(out.shape):
+                l = list(idx)
+                l[n_axis] += 1
+                assert out[idx] == x[tuple(l)] - x[idx], f"diff failed with {idx = }"
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise
 
 
 @pytest.mark.min_version("2024.12")
@@ -130,12 +145,19 @@ def test_diff_append_prepend(x, data):
     prepend_shape[n_axis] = prepend_axis_len
     prepend = data.draw(hh.arrays(dtype=x.dtype, shape=tuple(prepend_shape)), label="prepend")
 
-    out = xp.diff(x, **axis_kw, n=n, append=append, prepend=prepend)
+    repro_snippet = ph.format_snippet(
+        f"xp.diff({x!r}, **axis_kw, n={n!r}, append={append!r}, prepend={prepend!r}) with {axis_kw = }"
+    )
+    try:
+        out = xp.diff(x, **axis_kw, n=n, append=append, prepend=prepend)
 
-    in_1 = xp.concat((prepend, x, append), **axis_kw)
-    out_1 = xp.diff(in_1, **axis_kw, n=n)
+        in_1 = xp.concat((prepend, x, append), **axis_kw)
+        out_1 = xp.diff(in_1, **axis_kw, n=n)
 
-    assert out.shape == out_1.shape
-    for idx in sh.ndindex(out.shape):
-        assert out[idx] == out_1[idx], f"{idx = }"
+        assert out.shape == out_1.shape
+        for idx in sh.ndindex(out.shape):
+            assert out[idx] == out_1[idx], f"{idx = }"
+    except Exception as exc:
+        exc.add_note(repro_snippet)
+        raise