Factor out device checks into a helper function

Author: asmeurer

array_api_strict/_creation_functions.py

@@ -28,6 +28,14 @@ def _supports_buffer_protocol(obj):
         return False
     return True
 
+def _check_device(device):
+    # _array_object imports in this file are inside the functions to avoid
+    # circular imports
+    from ._array_object import CPU_DEVICE
+
+    if device not in [CPU_DEVICE, None]:
+        raise ValueError(f"Unsupported device {device!r}")
+
 def asarray(
     obj: Union[
         Array,
@@ -48,16 +56,13 @@ def asarray(
 
     See its docstring for more information.
     """
-    # _array_object imports in this file are inside the functions to avoid
-    # circular imports
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
     _np_dtype = None
     if dtype is not None:
         _np_dtype = dtype._np_dtype
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
 
     if np.__version__[0] < '2':
         if copy is False:
@@ -106,11 +111,11 @@ def arange(
 
     See its docstring for more information.
     """
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
+
     if dtype is not None:
         dtype = dtype._np_dtype
     return Array._new(np.arange(start, stop=stop, step=step, dtype=dtype))
@@ -127,11 +132,11 @@ def empty(
 
     See its docstring for more information.
     """
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
+
     if dtype is not None:
         dtype = dtype._np_dtype
     return Array._new(np.empty(shape, dtype=dtype))
@@ -145,11 +150,11 @@ def empty_like(
 
     See its docstring for more information.
     """
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
+
     if dtype is not None:
         dtype = dtype._np_dtype
     return Array._new(np.empty_like(x._array, dtype=dtype))
@@ -197,11 +202,11 @@ def full(
 
     See its docstring for more information.
     """
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
+
     if isinstance(fill_value, Array) and fill_value.ndim == 0:
         fill_value = fill_value._array
     if dtype is not None:
@@ -227,11 +232,11 @@ def full_like(
 
     See its docstring for more information.
     """
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
+
     if dtype is not None:
         dtype = dtype._np_dtype
     res = np.full_like(x._array, fill_value, dtype=dtype)
@@ -257,11 +262,11 @@ def linspace(
 
     See its docstring for more information.
     """
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
+
     if dtype is not None:
         dtype = dtype._np_dtype
     return Array._new(np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint))
@@ -298,11 +303,11 @@ def ones(
 
     See its docstring for more information.
     """
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
+
     if dtype is not None:
         dtype = dtype._np_dtype
     return Array._new(np.ones(shape, dtype=dtype))
@@ -316,11 +321,11 @@ def ones_like(
 
     See its docstring for more information.
     """
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
+
     if dtype is not None:
         dtype = dtype._np_dtype
     return Array._new(np.ones_like(x._array, dtype=dtype))
@@ -365,11 +370,11 @@ def zeros(
 
     See its docstring for more information.
     """
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
+
     if dtype is not None:
         dtype = dtype._np_dtype
     return Array._new(np.zeros(shape, dtype=dtype))
@@ -383,11 +388,11 @@ def zeros_like(
 
     See its docstring for more information.
     """
-    from ._array_object import Array, CPU_DEVICE
+    from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in [CPU_DEVICE, None]:
-        raise ValueError(f"Unsupported device {device!r}")
+    _check_device(device)
+
     if dtype is not None:
         dtype = dtype._np_dtype
     return Array._new(np.zeros_like(x._array, dtype=dtype))

array_api_strict/_data_type_functions.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from ._array_object import Array, CPU_DEVICE
+from ._array_object import Array
+from ._creation_functions import _check_device
 from ._dtypes import (
     _DType,
     _all_dtypes,
@@ -33,8 +34,7 @@ def astype(
 ) -> Array:
     if device is not _default:
         if get_array_api_strict_flags()['api_version'] >= '2023.12':
-            if device not in [CPU_DEVICE, None]:
-                raise ValueError(f"Unsupported device {device!r}")
+            _check_device(device)
         else:
             raise TypeError("The device argument to astype requires the 2023.12 version of the array API")