Move the array API dtype categories into the top level

They are not an official part of the spec but are useful for various parts of
the implementation.

Author: asmeurer

numpy/array_api/_array_object.py

@@ -98,23 +98,14 @@ def _check_allowed_dtypes(self, other, dtype_category, op):
             if other is NotImplemented:
                 return other
         """
-        from ._dtypes import _result_type
-
-        _dtypes = {
-            'all': _all_dtypes,
-            'numeric': _numeric_dtypes,
-            'integer': _integer_dtypes,
-            'integer or boolean': _integer_or_boolean_dtypes,
-            'boolean': _boolean_dtypes,
-            'floating-point': _floating_dtypes,
-        }
-
-        if self.dtype not in _dtypes[dtype_category]:
+        from ._dtypes import _result_type, _dtype_categories
+
+        if self.dtype not in _dtype_categories[dtype_category]:
             raise TypeError(f'Only {dtype_category} dtypes are allowed in {op}')
         if isinstance(other, (int, float, bool)):
             other = self._promote_scalar(other)
         elif isinstance(other, Array):
-            if other.dtype not in _dtypes[dtype_category]:
+            if other.dtype not in _dtype_categories[dtype_category]:
                 raise TypeError(f'Only {dtype_category} dtypes are allowed in {op}')
         else:
             return NotImplemented

numpy/array_api/_dtypes.py

@@ -23,6 +23,16 @@
 _integer_or_boolean_dtypes = (bool, int8, int16, int32, int64, uint8, uint16, uint32, uint64)
 _numeric_dtypes = (float32, float64, int8, int16, int32, int64, uint8, uint16, uint32, uint64)
 
+_dtype_categories = {
+    'all': _all_dtypes,
+    'numeric': _numeric_dtypes,
+    'integer': _integer_dtypes,
+    'integer or boolean': _integer_or_boolean_dtypes,
+    'boolean': _boolean_dtypes,
+    'floating-point': _floating_dtypes,
+}
+
+
 # Note: the spec defines a restricted type promotion table compared to NumPy.
 # In particular, cross-kind promotions like integer + float or boolean +
 # integer are not allowed, even for functions that accept both kinds.

numpy/array_api/tests/test_elementwise_functions.py

@@ -4,9 +4,8 @@
 
 from .. import asarray, _elementwise_functions
 from .._elementwise_functions import bitwise_left_shift, bitwise_right_shift
-from .._dtypes import (_all_dtypes, _boolean_dtypes, _floating_dtypes,
-                       _integer_dtypes, _integer_or_boolean_dtypes,
-                       _numeric_dtypes)
+from .._dtypes import (_dtype_categories, _boolean_dtypes, _floating_dtypes,
+                       _integer_dtypes)
 
 def nargs(func):
     return len(getfullargspec(func).args)
@@ -75,15 +74,6 @@ def test_function_types():
         'trunc': 'numeric',
     }
 
-    _dtypes = {
-        'all': _all_dtypes,
-        'numeric': _numeric_dtypes,
-        'integer': _integer_dtypes,
-        'integer_or_boolean': _integer_or_boolean_dtypes,
-        'boolean': _boolean_dtypes,
-        'floating': _floating_dtypes,
-    }
-
     def _array_vals():
         for d in _integer_dtypes:
             yield asarray(1, dtype=d)
@@ -94,7 +84,7 @@ def _array_vals():
 
     for x in _array_vals():
         for func_name, types in elementwise_function_input_types.items():
-            dtypes = _dtypes[types]
+            dtypes = _dtype_categories[types]
             func = getattr(_elementwise_functions, func_name)
             if nargs(func) == 2:
                 for y in _array_vals():