Default expected and out_name in ph.assert_dtype()

Author: honno

array_api_tests/pytest_helpers.py

@@ -1,5 +1,5 @@
 from inspect import getfullargspec
-from typing import Tuple
+from typing import Optional, Tuple
 
 from . import dtype_helpers as dh
 from . import function_stubs
@@ -43,12 +43,15 @@ def nargs(func_name):
 def assert_dtype(
     func_name: str,
     in_dtypes: Tuple[DataType, ...],
-    out_name: str,
     out_dtype: DataType,
-    expected: DataType
+    expected: Optional[DataType] = None,
+    *,
+    out_name: str = "out.dtype",
 ):
     f_in_dtypes = dh.fmt_types(in_dtypes)
     f_out_dtype = dh.dtype_to_name[out_dtype]
+    if expected is None:
+        expected = dh.result_type(*in_dtypes)
     f_expected = dh.dtype_to_name[expected]
     msg = (
         f"{out_name}={f_out_dtype}, but should be {f_expected} "

array_api_tests/test_linalg.py

@@ -271,13 +271,7 @@ def test_matmul(x1, x2):
     else:
         res = _array_module.matmul(x1, x2)
 
-    ph.assert_dtype(
-        "matmul",
-        (x1.dtype, x2.dtype),
-        "out.dtype",
-        res.dtype,
-        dh.promotion_table[x1.dtype, x2.dtype],
-    )
+    ph.assert_dtype("matmul", (x1.dtype, x2.dtype), res.dtype)
 
     if len(x1.shape) == len(x2.shape) == 1:
         assert res.shape == ()

array_api_tests/test_type_promotion.py

@@ -20,7 +20,7 @@
 @given(hh.mutually_promotable_dtypes(None))
 def test_result_type(dtypes):
     out = xp.result_type(*dtypes)
-    ph.assert_dtype('result_type', dtypes, 'out', out, dh.result_type(*dtypes))
+    ph.assert_dtype('result_type', dtypes, out, out_name='out')
 
 
 @given(
@@ -34,9 +34,8 @@ def test_meshgrid(dtypes, data):
         x = data.draw(xps.arrays(dtype=dtype, shape=shape), label=f'x{i}')
         arrays.append(x)
     out = xp.meshgrid(*arrays)
-    expected = dh.result_type(*dtypes)
     for i, x in enumerate(out):
-        ph.assert_dtype('meshgrid', dtypes, f'out[{i}].dtype', x.dtype, expected)
+        ph.assert_dtype('meshgrid', dtypes, x.dtype, out_name=f'out[{i}].dtype')
 
 
 @given(
@@ -50,7 +49,7 @@ def test_concat(shape, dtypes, data):
         x = data.draw(xps.arrays(dtype=dtype, shape=shape), label=f'x{i}')
         arrays.append(x)
     out = xp.concat(arrays)
-    ph.assert_dtype('concat', dtypes, 'out.dtype', out.dtype, dh.result_type(*dtypes))
+    ph.assert_dtype('concat', dtypes, out.dtype)
 
 
 @given(
@@ -64,7 +63,7 @@ def test_stack(shape, dtypes, data):
         x = data.draw(xps.arrays(dtype=dtype, shape=shape), label=f'x{i}')
         arrays.append(x)
     out = xp.stack(arrays)
-    ph.assert_dtype('stack', dtypes, 'out.dtype', out.dtype, dh.result_type(*dtypes))
+    ph.assert_dtype('stack', dtypes, out.dtype)
 
 
 bitwise_shift_funcs = [
@@ -150,7 +149,7 @@ def test_func_promotion(func_name, in_dtypes, out_dtype, data):
             out = func(*arrays)
         except OverflowError:
             reject()
-    ph.assert_dtype(func_name, in_dtypes, 'out.dtype', out.dtype, out_dtype)
+    ph.assert_dtype(func_name, in_dtypes, out.dtype, out_dtype)
 
 
 promotion_params: List[Param[Tuple[DataType, DataType], DataType]] = []
@@ -170,7 +169,7 @@ def test_where(in_dtypes, out_dtype, shapes, data):
     x2 = data.draw(xps.arrays(dtype=in_dtypes[1], shape=shapes[1]), label='x2')
     cond = data.draw(xps.arrays(dtype=xp.bool, shape=shapes[2]), label='condition')
     out = xp.where(cond, x1, x2)
-    ph.assert_dtype('where', in_dtypes, 'out.dtype', out.dtype, out_dtype)
+    ph.assert_dtype('where', in_dtypes, out.dtype, out_dtype)
 
 
 numeric_promotion_params = promotion_params[1:]
@@ -182,7 +181,7 @@ def test_tensordot(in_dtypes, out_dtype, shapes, data):
     x1 = data.draw(xps.arrays(dtype=in_dtypes[0], shape=shapes[0]), label='x1')
     x2 = data.draw(xps.arrays(dtype=in_dtypes[1], shape=shapes[1]), label='x2')
     out = xp.tensordot(x1, x2)
-    ph.assert_dtype('tensordot', in_dtypes, 'out.dtype', out.dtype, out_dtype)
+    ph.assert_dtype('tensordot', in_dtypes, out.dtype, out_dtype)
 
 
 @pytest.mark.parametrize('in_dtypes, out_dtype', numeric_promotion_params)
@@ -191,7 +190,7 @@ def test_vecdot(in_dtypes, out_dtype, shapes, data):
     x1 = data.draw(xps.arrays(dtype=in_dtypes[0], shape=shapes[0]), label='x1')
     x2 = data.draw(xps.arrays(dtype=in_dtypes[1], shape=shapes[1]), label='x2')
     out = xp.vecdot(x1, x2)
-    ph.assert_dtype('vecdot', in_dtypes, 'out.dtype', out.dtype, out_dtype)
+    ph.assert_dtype('vecdot', in_dtypes, out.dtype, out_dtype)
 
 
 op_params: List[Param[str, str, Tuple[DataType, ...], DataType]] = []
@@ -259,7 +258,7 @@ def test_op_promotion(op, expr, in_dtypes, out_dtype, data):
             out = eval(expr, locals_)
         except OverflowError:
             reject()
-    ph.assert_dtype(op, in_dtypes, 'out.dtype', out.dtype, out_dtype)
+    ph.assert_dtype(op, in_dtypes, out.dtype, out_dtype)
 
 
 inplace_params: List[Param[str, str, Tuple[DataType, ...], DataType]] = []
@@ -300,7 +299,7 @@ def test_inplace_op_promotion(op, expr, in_dtypes, out_dtype, shapes, data):
     except OverflowError:
         reject()
     x1 = locals_['x1']
-    ph.assert_dtype(op, in_dtypes, 'x1.dtype', x1.dtype, out_dtype)
+    ph.assert_dtype(op, in_dtypes, x1.dtype, out_dtype, out_name='x1.dtype')
 
 
 op_scalar_params: List[Param[str, str, DataType, ScalarType, DataType]] = []
@@ -334,7 +333,7 @@ def test_op_scalar_promotion(op, expr, in_dtype, in_stype, out_dtype, data):
         out = eval(expr, {'x': x, 's': s})
     except OverflowError:
         reject()
-    ph.assert_dtype(op, (in_dtype, in_stype), 'out.dtype', out.dtype, out_dtype)
+    ph.assert_dtype(op, (in_dtype, in_stype), out.dtype, out_dtype)
 
 
 inplace_scalar_params: List[Param[str, str, DataType, ScalarType]] = []
@@ -369,7 +368,7 @@ def test_inplace_op_scalar_promotion(op, expr, dtype, in_stype, data):
         reject()
     x = locals_['x']
     assert x.dtype == dtype, f'{x.dtype=!s}, but should be {dtype}'
-    ph.assert_dtype(op, (dtype, in_stype), 'x.dtype', x.dtype, dtype)
+    ph.assert_dtype(op, (dtype, in_stype), x.dtype, dtype, out_name='x.dtype')
 
 
 if __name__ == '__main__':