Use dtype object instead of name for dtype helpers

Author: honno

array_api_tests/array_helpers.py

@@ -9,7 +9,7 @@
                             _numeric_dtypes, _boolean_dtypes, _dtypes,
                             asarray)
 from . import _array_module
-from .dtype_helpers import dtype_mapping, promotion_table
+from .dtype_helpers import promotion_table
 
 # These are exported here so that they can be included in the special cases
 # tests from this file.
@@ -371,14 +371,10 @@ def promote_dtypes(dtype1, dtype2):
     Special case of result_type() which uses the exact type promotion table
     from the spec.
     """
-    # Equivalent to this, but some libraries may not work properly with using
-    # dtype objects as dict keys
-    #
-    # d1, d2 = reverse_dtype_mapping[dtype1], reverse_dtype_mapping[dtype2]
-
-    d1 = [i for i in dtype_mapping if dtype_mapping[i] == dtype1][0]
-    d2 = [i for i in dtype_mapping if dtype_mapping[i] == dtype2][0]
-
-    if (d1, d2) not in promotion_table:
-        raise ValueError(f"{d1} and {d2} are not type promotable according to the spec (this may indicate a bug in the test suite).")
-    return dtype_mapping[promotion_table[d1, d2]]
+    try:
+        return promotion_table[(dtype2, dtype2)]
+    except KeyError as e:
+        raise ValueError(
+            f"{dtype1} and {dtype2} are not type promotable according to the spec"
+            f"(this may indicate a bug in the test suite)."
+        ) from e

array_api_tests/dtype_helpers.py

@@ -1,7 +1,7 @@
 from . import _array_module as xp
 
+
 __all__ = [
-    "dtype_mapping",
     "promotion_table",
     "dtype_nbits",
     "dtype_signed",
@@ -14,21 +14,6 @@
     "operators_to_functions",
 ]
 
-dtype_mapping = {
-    'int8': xp.int8,
-    'int16': xp.int16,
-    'int32': xp.int32,
-    'int64': xp.int64,
-    'uint8': xp.uint8,
-    'uint16': xp.uint16,
-    'uint32': xp.uint32,
-    'uint64': xp.uint64,
-    'float32': xp.float32,
-    'float64': xp.float64,
-    'bool': xp.bool,
-}
-
-reverse_dtype_mapping = {v: k for k, v in dtype_mapping.items()}
 
 def dtype_nbits(dtype):
     if dtype == xp.int8:
@@ -54,79 +39,87 @@ def dtype_nbits(dtype):
     else:
         raise ValueError(f"dtype_nbits is not defined for {dtype}")
 
+
 def dtype_signed(dtype):
     if dtype in [xp.int8, xp.int16, xp.int32, xp.int64]:
         return True
     elif dtype in [xp.uint8, xp.uint16, xp.uint32, xp.uint64]:
         return False
     raise ValueError("dtype_signed is only defined for integer dtypes")
 
+
 signed_integer_promotion_table = {
-    ('int8', 'int8'): 'int8',
-    ('int8', 'int16'): 'int16',
-    ('int8', 'int32'): 'int32',
-    ('int8', 'int64'): 'int64',
-    ('int16', 'int8'): 'int16',
-    ('int16', 'int16'): 'int16',
-    ('int16', 'int32'): 'int32',
-    ('int16', 'int64'): 'int64',
-    ('int32', 'int8'): 'int32',
-    ('int32', 'int16'): 'int32',
-    ('int32', 'int32'): 'int32',
-    ('int32', 'int64'): 'int64',
-    ('int64', 'int8'): 'int64',
-    ('int64', 'int16'): 'int64',
-    ('int64', 'int32'): 'int64',
-    ('int64', 'int64'): 'int64',
+    (xp.int8, xp.int8): xp.int8,
+    (xp.int8, xp.int16): xp.int16,
+    (xp.int8, xp.int32): xp.int32,
+    (xp.int8, xp.int64): xp.int64,
+    (xp.int16, xp.int8): xp.int16,
+    (xp.int16, xp.int16): xp.int16,
+    (xp.int16, xp.int32): xp.int32,
+    (xp.int16, xp.int64): xp.int64,
+    (xp.int32, xp.int8): xp.int32,
+    (xp.int32, xp.int16): xp.int32,
+    (xp.int32, xp.int32): xp.int32,
+    (xp.int32, xp.int64): xp.int64,
+    (xp.int64, xp.int8): xp.int64,
+    (xp.int64, xp.int16): xp.int64,
+    (xp.int64, xp.int32): xp.int64,
+    (xp.int64, xp.int64): xp.int64,
 }
 
+
 unsigned_integer_promotion_table = {
-    ('uint8', 'uint8'): 'uint8',
-    ('uint8', 'uint16'): 'uint16',
-    ('uint8', 'uint32'): 'uint32',
-    ('uint8', 'uint64'): 'uint64',
-    ('uint16', 'uint8'): 'uint16',
-    ('uint16', 'uint16'): 'uint16',
-    ('uint16', 'uint32'): 'uint32',
-    ('uint16', 'uint64'): 'uint64',
-    ('uint32', 'uint8'): 'uint32',
-    ('uint32', 'uint16'): 'uint32',
-    ('uint32', 'uint32'): 'uint32',
-    ('uint32', 'uint64'): 'uint64',
-    ('uint64', 'uint8'): 'uint64',
-    ('uint64', 'uint16'): 'uint64',
-    ('uint64', 'uint32'): 'uint64',
-    ('uint64', 'uint64'): 'uint64',
+    (xp.uint8, xp.uint8): xp.uint8,
+    (xp.uint8, xp.uint16): xp.uint16,
+    (xp.uint8, xp.uint32): xp.uint32,
+    (xp.uint8, xp.uint64): xp.uint64,
+    (xp.uint16, xp.uint8): xp.uint16,
+    (xp.uint16, xp.uint16): xp.uint16,
+    (xp.uint16, xp.uint32): xp.uint32,
+    (xp.uint16, xp.uint64): xp.uint64,
+    (xp.uint32, xp.uint8): xp.uint32,
+    (xp.uint32, xp.uint16): xp.uint32,
+    (xp.uint32, xp.uint32): xp.uint32,
+    (xp.uint32, xp.uint64): xp.uint64,
+    (xp.uint64, xp.uint8): xp.uint64,
+    (xp.uint64, xp.uint16): xp.uint64,
+    (xp.uint64, xp.uint32): xp.uint64,
+    (xp.uint64, xp.uint64): xp.uint64,
 }
 
+
 mixed_signed_unsigned_promotion_table = {
-    ('int8', 'uint8'): 'int16',
-    ('int8', 'uint16'): 'int32',
-    ('int8', 'uint32'): 'int64',
-    ('int16', 'uint8'): 'int16',
-    ('int16', 'uint16'): 'int32',
-    ('int16', 'uint32'): 'int64',
-    ('int32', 'uint8'): 'int32',
-    ('int32', 'uint16'): 'int32',
-    ('int32', 'uint32'): 'int64',
-    ('int64', 'uint8'): 'int64',
-    ('int64', 'uint16'): 'int64',
-    ('int64', 'uint32'): 'int64',
+    (xp.int8, xp.uint8): xp.int16,
+    (xp.int8, xp.uint16): xp.int32,
+    (xp.int8, xp.uint32): xp.int64,
+    (xp.int16, xp.uint8): xp.int16,
+    (xp.int16, xp.uint16): xp.int32,
+    (xp.int16, xp.uint32): xp.int64,
+    (xp.int32, xp.uint8): xp.int32,
+    (xp.int32, xp.uint16): xp.int32,
+    (xp.int32, xp.uint32): xp.int64,
+    (xp.int64, xp.uint8): xp.int64,
+    (xp.int64, xp.uint16): xp.int64,
+    (xp.int64, xp.uint32): xp.int64,
 }
 
+
 flipped_mixed_signed_unsigned_promotion_table = {(u, i): p for (i, u), p in mixed_signed_unsigned_promotion_table.items()}
 
+
 float_promotion_table = {
-    ('float32', 'float32'): 'float32',
-    ('float32', 'float64'): 'float64',
-    ('float64', 'float32'): 'float64',
-    ('float64', 'float64'): 'float64',
+    (xp.float32, xp.float32): xp.float32,
+    (xp.float32, xp.float64): xp.float64,
+    (xp.float64, xp.float32): xp.float64,
+    (xp.float64, xp.float64): xp.float64,
 }
 
+
 boolean_promotion_table = {
-    ('bool', 'bool'): 'bool',
+    (xp.bool, xp.bool): xp.bool,
 }
 
+
 promotion_table = {
     **signed_integer_promotion_table,
     **unsigned_integer_promotion_table,
@@ -136,6 +129,7 @@ def dtype_signed(dtype):
     **boolean_promotion_table,
 }
 
+
 input_types = {
     'any': sorted(set(promotion_table.values())),
     'boolean': sorted(set(boolean_promotion_table.values())),
@@ -150,21 +144,23 @@ def dtype_signed(dtype):
                            **unsigned_integer_promotion_table}.values())),
 }
 
+
 dtypes_to_scalars = {
-    'bool': [bool],
-    'int8': [int],
-    'int16': [int],
-    'int32': [int],
-    'int64': [int],
+    xp.bool: [bool],
+    xp.int8: [int],
+    xp.int16: [int],
+    xp.int32: [int],
+    xp.int64: [int],
     # Note: unsigned int dtypes only correspond to positive integers
-    'uint8': [int],
-    'uint16': [int],
-    'uint32': [int],
-    'uint64': [int],
-    'float32': [int, float],
-    'float64': [int, float],
+    xp.uint8: [int],
+    xp.uint16: [int],
+    xp.uint32: [int],
+    xp.uint64: [int],
+    xp.float32: [int, float],
+    xp.float64: [int, float],
 }
 
+
 elementwise_function_input_types = {
     'abs': 'numeric',
     'acos': 'floating',
@@ -224,6 +220,7 @@ def dtype_signed(dtype):
     'trunc': 'numeric',
 }
 
+
 elementwise_function_output_types = {
     'abs': 'promoted',
     'acos': 'promoted',
@@ -283,6 +280,7 @@ def dtype_signed(dtype):
     'trunc': 'promoted',
 }
 
+
 binary_operators = {
     '__add__': '+',
     '__and__': '&',
@@ -305,6 +303,7 @@ def dtype_signed(dtype):
     '__xor__': '^',
 }
 
+
 unary_operators = {
     '__abs__': 'abs()',
     '__invert__': '~',

array_api_tests/hypothesis_helpers.py

@@ -14,7 +14,7 @@
                             boolean_dtype_objects,
                             integer_or_boolean_dtype_objects, dtype_objects)
 from ._array_module import full, float32, float64, bool as bool_dtype, _UndefinedStub
-from .dtype_helpers import dtype_mapping, promotion_table
+from .dtype_helpers import promotion_table
 from . import _array_module
 from . import _array_module as xp
 
@@ -54,9 +54,8 @@
 sorted_table = sorted(
     sorted_table, key=lambda ij: -1 if ij[0] == ij[1] else sorted_table.index(ij)
 )
-dtype_pairs = [(dtype_mapping[i], dtype_mapping[j]) for i, j in sorted_table]
 if FILTER_UNDEFINED_DTYPES:
-    dtype_pairs = [(i, j) for i, j in dtype_pairs
+    sorted_table = [(i, j) for i, j in sorted_table
                     if not isinstance(i, _UndefinedStub)
                     and not isinstance(j, _UndefinedStub)]
 
@@ -70,7 +69,7 @@ def mutually_promotable_dtypes(dtype_objects=dtype_objects):
     # pairs (XXX: Can we redesign the strategies so that they can prefer
     # shrinking dtypes over values?)
     return sampled_from(
-        [(i, j) for i, j in dtype_pairs if i in dtype_objects and j in dtype_objects]
+        [(i, j) for i, j in sorted_table if i in dtype_objects and j in dtype_objects]
     )
 
 shared_mutually_promotable_dtype_pairs = shared(

array_api_tests/test_broadcasting.py

@@ -10,8 +10,7 @@
 from .hypothesis_helpers import shapes, FILTER_UNDEFINED_DTYPES
 from .pytest_helpers import raises, doesnt_raise, nargs
 
-from .dtype_helpers import (elementwise_function_input_types,
-                                  input_types, dtype_mapping)
+from .dtype_helpers import elementwise_function_input_types, input_types
 from .function_stubs import elementwise_functions
 from . import _array_module
 from ._array_module import ones, _UndefinedStub
@@ -111,14 +110,14 @@ def test_broadcast_shapes_explicit_spec():
 @pytest.mark.parametrize('func_name', [i for i in
                                        elementwise_functions.__all__ if
                                        nargs(i) > 1])
-@given(shape1=shapes, shape2=shapes, dtype=data())
-def test_broadcasting_hypothesis(func_name, shape1, shape2, dtype):
+@given(shape1=shapes, shape2=shapes, data=data())
+def test_broadcasting_hypothesis(func_name, shape1, shape2, data):
     # Internal consistency checks
     assert nargs(func_name) == 2
 
-    dtype = dtype_mapping[dtype.draw(sampled_from(input_types[elementwise_function_input_types[func_name]]))]
-    if FILTER_UNDEFINED_DTYPES and isinstance(dtype, _UndefinedStub):
-        assume(False)
+    dtype = data.draw(sampled_from(input_types[elementwise_function_input_types[func_name]]))
+    if FILTER_UNDEFINED_DTYPES:
+        assume(not isinstance(dtype, _UndefinedStub))
     func = getattr(_array_module, func_name)
 
     if isinstance(func, _array_module._UndefinedStub):