Manipulation and statistical tests, smoke all untested functions

array_api_tests/meta/test_utils.py

@@ -1,7 +1,10 @@
 import pytest
 
-from ..test_signatures import extension_module
+from .. import array_helpers as ah
 from ..test_creation_functions import frange
+from ..test_manipulation_functions import axis_ndindex
+from ..test_signatures import extension_module
+from ..test_statistical_functions import axes_ndindex
 
 
 def test_extension_module_is_extension():
@@ -24,3 +27,45 @@ def test_extension_func_is_not_extension():
 def test_frange(r, size, elements):
     assert len(r) == size
     assert list(r) == elements
+
+
+@pytest.mark.parametrize(
+    "shape, expected",
+    [((), [()])],
+)
+def test_ndindex(shape, expected):
+    assert list(ah.ndindex(shape)) == expected
+
+
+@pytest.mark.parametrize(
+    "shape, axis, expected",
+    [
+        ((1,), 0, [(slice(None, None),)]),
+        ((1, 2), 0, [(slice(None, None), slice(None, None))]),
+        (
+            (2, 4),
+            1,
+            [(0, slice(None, None)), (1, slice(None, None))],
+        ),
+    ],
+)
+def test_axis_ndindex(shape, axis, expected):
+    assert list(axis_ndindex(shape, axis)) == expected
+
+
+@pytest.mark.parametrize(
+    "shape, axes, expected",
+    [
+        ((), (), [((),)]),
+        (
+            (2, 2),
+            (0,),
+            [
+                ((0, 0), (1, 0)),
+                ((0, 1), (1, 1)),
+            ],
+        ),
+    ],
+)
+def test_axes_ndindex(shape, axes, expected):
+    assert list(axes_ndindex(shape, axes)) == expected

array_api_tests/pytest_helpers.py

@@ -67,12 +67,14 @@ def fmt_kw(kw: Dict[str, Any]) -> str:
 
 def assert_dtype(
     func_name: str,
-    in_dtypes: Tuple[DataType, ...],
+    in_dtypes: Union[DataType, Tuple[DataType, ...]],
     out_dtype: DataType,
     expected: Optional[DataType] = None,
     *,
     repr_name: str = "out.dtype",
 ):
+    if not isinstance(in_dtypes, tuple):
+        in_dtypes = (in_dtypes,)
     f_in_dtypes = dh.fmt_types(in_dtypes)
     f_out_dtype = dh.dtype_to_name[out_dtype]
     if expected is None:
@@ -149,9 +151,7 @@ def assert_result_shape(
     f_sig = f" {f_in_shapes} "
     if kw:
         f_sig += f", {fmt_kw(kw)}"
-    msg = (
-        f"{repr_name}={out_shape}, but should be {expected} [{func_name}({f_sig})]"
-    )
+    msg = f"{repr_name}={out_shape}, but should be {expected} [{func_name}({f_sig})]"
     assert out_shape == expected, msg