Add device flag to astype in 2023.12

Also clean up imports in test_data_type_functions.py

Author: asmeurer

array_api_strict/_data_type_functions.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ._array_object import Array
+from ._array_object import Array, CPU_DEVICE
 from ._dtypes import (
     _DType,
     _all_dtypes,
@@ -13,19 +13,31 @@
     _numeric_dtypes,
     _result_type,
 )
+from ._flags import get_array_api_strict_flags
 
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import List, Tuple, Union
-    from ._typing import Dtype
+    from typing import List, Tuple, Union, Optional
+    from ._typing import Dtype, Device
 
 import numpy as np
 
+# Use to emulate the asarray(device) argument not existing in 2022.12
+_default = object()
 
 # Note: astype is a function, not an array method as in NumPy.
-def astype(x: Array, dtype: Dtype, /, *, copy: bool = True) -> Array:
+def astype(
+    x: Array, dtype: Dtype, /, *, copy: bool = True, device: Optional[Device] = _default
+) -> 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}")
+        else:
+            raise TypeError("The device argument to astype requires the 2023.12 version of the array API")
+
     if not copy and dtype == x.dtype:
         return x
     return Array._new(x._array.astype(dtype=dtype._np_dtype, copy=copy))

array_api_strict/tests/test_data_type_functions.py

@@ -3,38 +3,68 @@
 import pytest
 
 from numpy.testing import assert_raises
-import array_api_strict as xp
 import numpy as np
 
+from .._creation_functions import asarray
+from .._data_type_functions import astype, can_cast, isdtype
+from .._dtypes import (
+    bool, int8, int16, uint8, float64,
+)
+from .._flags import set_array_api_strict_flags
+
+
 @pytest.mark.parametrize(
     "from_, to, expected",
     [
-        (xp.int8, xp.int16, True),
-        (xp.int16, xp.int8, False),
-        (xp.bool, xp.int8, False),
-        (xp.asarray(0, dtype=xp.uint8), xp.int8, False),
+        (int8, int16, True),
+        (int16, int8, False),
+        (bool, int8, False),
+        (asarray(0, dtype=uint8), int8, False),
     ],
 )
 def test_can_cast(from_, to, expected):
     """
     can_cast() returns correct result
     """
-    assert xp.can_cast(from_, to) == expected
+    assert can_cast(from_, to) == expected
 
 def test_isdtype_strictness():
-    assert_raises(TypeError, lambda: xp.isdtype(xp.float64, 64))
-    assert_raises(ValueError, lambda: xp.isdtype(xp.float64, 'f8'))
+    assert_raises(TypeError, lambda: isdtype(float64, 64))
+    assert_raises(ValueError, lambda: isdtype(float64, 'f8'))
 
-    assert_raises(TypeError, lambda: xp.isdtype(xp.float64, (('integral',),)))
+    assert_raises(TypeError, lambda: isdtype(float64, (('integral',),)))
     with assert_raises(TypeError), warnings.catch_warnings(record=True) as w:
         warnings.simplefilter("always")
-        xp.isdtype(xp.float64, np.object_)
+        isdtype(float64, np.object_)
         assert len(w) == 1
         assert issubclass(w[-1].category, UserWarning)
 
-    assert_raises(TypeError, lambda: xp.isdtype(xp.float64, None))
+    assert_raises(TypeError, lambda: isdtype(float64, None))
     with assert_raises(TypeError), warnings.catch_warnings(record=True) as w:
         warnings.simplefilter("always")
-        xp.isdtype(xp.float64, np.float64)
+        isdtype(float64, np.float64)
         assert len(w) == 1
         assert issubclass(w[-1].category, UserWarning)
+
+
+@pytest.mark.parametrize("api_version", ['2021.12', '2022.12', '2023.12'])
+def astype_device(api_version):
+    if api_version != '2022.12':
+        with pytest.warns(UserWarning):
+            set_array_api_strict_flags(api_version=api_version)
+    else:
+        set_array_api_strict_flags(api_version=api_version)
+
+    a = asarray([1, 2, 3], dtype=int8)
+    # Never an error
+    astype(a, int16)
+
+    # Always an error
+    astype(a, int16, device="cpu")
+
+    if api_version >= '2023.12':
+        astype(a, int8, device=None)
+        astype(a, int8, device=a.device)
+    else:
+        pytest.raises(TypeError, lambda: astype(a, int8, device=None))
+        pytest.raises(TypeError, lambda: astype(a, int8, device=a.device))