Test more valid argument signatures in test_arange

Author: honno

array_api_tests/pytest_helpers.py

@@ -11,6 +11,7 @@
     "raises",
     "doesnt_raise",
     "nargs",
+    "fmt_kw",
     "assert_dtype",
     "assert_kw_dtype",
     "assert_default_float",
@@ -58,7 +59,7 @@ def nargs(func_name):
     return len(getfullargspec(getattr(function_stubs, func_name)).args)
 
 
-def _fmt_kw(kw: Dict[str, Any]) -> str:
+def fmt_kw(kw: Dict[str, Any]) -> str:
     return ", ".join(f"{k}={v}" for k, v in kw.items())
 
 
@@ -120,15 +121,15 @@ def assert_shape(
     if isinstance(expected, int):
         expected = (expected,)
     msg = (
-        f"out.shape={out_shape}, but should be {expected} [{func_name}({_fmt_kw(kw)})]"
+        f"out.shape={out_shape}, but should be {expected} [{func_name}({fmt_kw(kw)})]"
     )
     assert out_shape == expected, msg
 
 
 def assert_fill(
     func_name: str, fill_value: Scalar, dtype: DataType, out: Array, /, **kw
 ):
-    msg = f"out not filled with {fill_value} [{func_name}({_fmt_kw(kw)})]\n{out=}"
+    msg = f"out not filled with {fill_value} [{func_name}({fmt_kw(kw)})]\n{out=}"
     if math.isnan(fill_value):
         assert ah.all(ah.isnan(out)), msg
     else:

array_api_tests/test_creation_functions.py

@@ -131,15 +131,16 @@ def test_arange(dtype, data):
         size <= hh.MAX_ARRAY_SIZE
     ), f"{size=} should be no more than {hh.MAX_ARRAY_SIZE}"  # sanity check
 
-    kw = data.draw(
-        hh.specified_kwargs(
-            hh.KVD("stop", stop, None),
-            hh.KVD("step", step, None),
-            hh.KVD("dtype", dtype, None),
-        ),
-        label="kw",
-    )
-    out = xp.arange(start, **kw)
+    args_samples = [(start, stop), (start, stop, step)]
+    if stop is None:
+        args_samples.insert(0, (start,))
+    args = data.draw(st.sampled_from(args_samples), label="args")
+    kvds = [hh.KVD("dtype", dtype, None)]
+    if len(args) != 3:
+        kvds.insert(0, hh.KVD("step", step, 1))
+    kwargs = data.draw(hh.specified_kwargs(*kvds), label="kwargs")
+
+    out = xp.arange(*args, **kwargs)
 
     if dtype is None:
         if all_int:
@@ -148,8 +149,11 @@ def test_arange(dtype, data):
             ph.assert_default_float("arange", out.dtype)
     else:
         ph.assert_dtype("arange", (out.dtype,), dtype)
-    assert out.ndim == 1, f"{out.ndim=}, but should be 1 [linspace()]"
-    f_func = f"[arange({start}, {stop}, {step})]"
+    f_sig = ", ".join(str(n) for n in args)
+    if len(kwargs) > 0:
+        f_sig += f", {ph.fmt_kw(kwargs)}"
+    f_func = f"[arange({f_sig})]"
+    assert out.ndim == 1, f"{out.ndim=}, but should be 1 [{f_func}]"
     # We check size is roughly as expected to avoid edge cases e.g.
     #
     #     >>> xp.arange(2, step=0.333333333333333)