Convert dtype_signed and dtype_nbits to dicts

Author: honno

array_api_tests/dtype_helpers.py

@@ -15,37 +15,18 @@
 ]
 
 
-def dtype_nbits(dtype):
-    if dtype == xp.int8:
-        return 8
-    elif dtype == xp.int16:
-        return 16
-    elif dtype == xp.int32:
-        return 32
-    elif dtype == xp.int64:
-        return 64
-    elif dtype == xp.uint8:
-        return 8
-    elif dtype == xp.uint16:
-        return 16
-    elif dtype == xp.uint32:
-        return 32
-    elif dtype == xp.uint64:
-        return 64
-    elif dtype == xp.float32:
-        return 32
-    elif dtype == xp.float64:
-        return 64
-    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")
+dtype_nbits = {
+    **{d: 8 for d in [xp.int8, xp.uint8]},
+    **{d: 16 for d in [xp.int16, xp.uint16]},
+    **{d: 32 for d in [xp.int32, xp.uint32, xp.float32]},
+    **{d: 64 for d in [xp.int64, xp.uint64, xp.float64]},
+}
+
+
+dtype_signed = {
+    **{d: True for d in [xp.int8, xp.int16, xp.int32, xp.int64]},
+    **{d: False for d in [xp.uint8, xp.uint16, xp.uint32, xp.uint64]},
+}
 
 
 signed_integer_promotion_table = {

array_api_tests/meta_tests/test_array_helpers.py

@@ -25,8 +25,8 @@ def test_notequal():
 
 @given(integers(), integer_dtypes)
 def test_int_to_dtype(x, dtype):
-    n = dtype_nbits(dtype)
-    signed = dtype_signed(dtype)
+    n = dtype_nbits[dtype]
+    signed = dtype_signed[dtype]
     try:
         d = xp.asarray(x, dtype=dtype)
     except OverflowError:

array_api_tests/test_elementwise_functions.py

@@ -223,7 +223,7 @@ def test_bitwise_and(args):
         x = int(x1)
         y = int(x2)
         res = int(a)
-        ans = int_to_dtype(x & y, dtype_nbits(a.dtype), dtype_signed(a.dtype))
+        ans = int_to_dtype(x & y, dtype_nbits[a.dtype], dtype_signed[a.dtype])
         assert ans == res
 
 @given(two_integer_dtypes.flatmap(lambda i: two_array_scalars(*i)))
@@ -240,12 +240,12 @@ def test_bitwise_left_shift(args):
         raise RuntimeError("Error: test_bitwise_left_shift needs to be updated for nonscalar array inputs")
     x = int(x1)
     y = int(x2)
-    if y >= dtype_nbits(a.dtype):
+    if y >= dtype_nbits[a.dtype]:
         # Avoid shifting very large y in Python ints
         ans = 0
     else:
         ans = x << y
-    ans = int_to_dtype(ans, dtype_nbits(a.dtype), dtype_signed(a.dtype))
+    ans = int_to_dtype(ans, dtype_nbits[a.dtype], dtype_signed[a.dtype])
     res = int(a)
     assert ans == res
 
@@ -263,7 +263,7 @@ def test_bitwise_invert(x):
     else:
         x = int(x)
         res = int(a)
-        ans = int_to_dtype(~x, dtype_nbits(a.dtype), dtype_signed(a.dtype))
+        ans = int_to_dtype(~x, dtype_nbits[a.dtype], dtype_signed[a.dtype])
         assert ans == res
 
 @given(two_integer_or_boolean_dtypes.flatmap(lambda i: two_array_scalars(*i)))
@@ -284,7 +284,7 @@ def test_bitwise_or(args):
         x = int(x1)
         y = int(x2)
         res = int(a)
-        ans = int_to_dtype(x | y, dtype_nbits(a.dtype), dtype_signed(a.dtype))
+        ans = int_to_dtype(x | y, dtype_nbits[a.dtype], dtype_signed[a.dtype])
         assert ans == res
 
 @given(two_integer_dtypes.flatmap(lambda i: two_array_scalars(*i)))
@@ -301,7 +301,7 @@ def test_bitwise_right_shift(args):
         raise RuntimeError("Error: test_bitwise_right_shift needs to be updated for nonscalar array inputs")
     x = int(x1)
     y = int(x2)
-    ans = int_to_dtype(x >> y, dtype_nbits(a.dtype), dtype_signed(a.dtype))
+    ans = int_to_dtype(x >> y, dtype_nbits[a.dtype], dtype_signed[a.dtype])
     res = int(a)
     assert ans == res
 
@@ -323,7 +323,7 @@ def test_bitwise_xor(args):
         x = int(x1)
         y = int(x2)
         res = int(a)
-        ans = int_to_dtype(x ^ y, dtype_nbits(a.dtype), dtype_signed(a.dtype))
+        ans = int_to_dtype(x ^ y, dtype_nbits[a.dtype], dtype_signed[a.dtype])
         assert ans == res
 
 @given(numeric_scalars)