Construct inplace_op_to_symbol in dtype_helpers

Author: honno

array_api_tests/dtype_helpers.py

@@ -23,9 +23,10 @@
     'func_out_categories',
     'op_in_categories',
     'op_out_categories',
+    'op_to_func',
     'binary_op_to_symbol',
     'unary_op_to_symbol',
-    'op_to_func',
+    'inplace_op_to_symbol',
 ]
 
 
@@ -328,6 +329,16 @@ class MinMax(NamedTuple):
 
 op_in_categories = {}
 op_out_categories = {}
-for op_func, elwise_func in op_to_func.items():
-    op_in_categories[op_func] = func_in_categories[elwise_func]
-    op_out_categories[op_func] = func_out_categories[elwise_func]
+for op, elwise_func in op_to_func.items():
+    op_in_categories[op] = func_in_categories[elwise_func]
+    op_out_categories[op] = func_out_categories[elwise_func]
+
+
+inplace_op_to_symbol = {}
+for op, symbol in binary_op_to_symbol.items():
+    if op == '__matmul__' or op_out_categories[op] == 'bool':
+        continue
+    iop = f'__i{op[2:]}'
+    inplace_op_to_symbol[iop] = f'{symbol}='
+    op_in_categories[iop] = op_in_categories[op]
+    op_out_categories[iop] = op_out_categories[op]

array_api_tests/test_type_promotion.py

@@ -168,24 +168,23 @@ def test_operator_returns_array_with_correct_dtype(
 
 
 def gen_inplace_params() -> Iterator[Tuple[str, Tuple[DT, ...], DT, Callable]]:
-    for op, symbol in dh.binary_op_to_symbol.items():
-        if op == '__matmul__' or dh.op_out_categories[op] == 'bool':
+    for op, symbol in dh.inplace_op_to_symbol.items():
+        if op == '__imatmul__':
             continue
         in_category = dh.op_in_categories[op]
         valid_in_dtypes = dh.category_to_dtypes[in_category]
-        iop = f'__i{op[2:]}'
         for (in_dtype1, in_dtype2), promoted_dtype in dh.promotion_table.items():
             if (
                 in_dtype1 == promoted_dtype
                 and in_dtype1 in valid_in_dtypes
                 and in_dtype2 in valid_in_dtypes
             ):
                 yield pytest.param(
-                    f'x1 {symbol}= x2',
+                    f'x1 {symbol} x2',
                     (in_dtype1, in_dtype2),
                     promoted_dtype,
-                    filters[iop],
-                    id=f'{iop}({in_dtype1}, {in_dtype2}) -> {promoted_dtype}',
+                    filters[op],
+                    id=f'{op}({in_dtype1}, {in_dtype2}) -> {promoted_dtype}',
                 )
 
 
@@ -252,19 +251,18 @@ def test_binary_operator_promotes_python_scalars(
 
 
 def gen_inplace_scalar_params() -> Iterator[Tuple[str, DT, ScalarType, Callable]]:
-    for op, symbol in dh.binary_op_to_symbol.items():
-        if op == '__matmul__' or dh.op_out_categories[op] == 'bool':
+    for op, symbol in dh.inplace_op_to_symbol.items():
+        if op == '__imatmul__':
             continue
         in_category = dh.op_in_categories[op]
-        iop = f'__i{op[2:]}'
         for dtype in dh.category_to_dtypes[in_category]:
             for in_stype in dh.dtypes_to_scalars[dtype]:
                 yield pytest.param(
-                    f'x {symbol}= s',
+                    f'x {symbol} s',
                     dtype,
                     in_stype,
-                    filters[iop],
-                    id=f'{iop}({dtype}, {in_stype.__name__}) -> {dtype}',
+                    filters[op],
+                    id=f'{op}({dtype}, {in_stype.__name__}) -> {dtype}',
                 )