🏷️: add Device parametrization

Signed-off-by: nstarman <[email protected]>

Author: nstarman

src/array_api_typing/__init__.py

@@ -4,6 +4,7 @@
     "Array",
     "HasArrayNamespace",
     "HasDType",
+    "HasDevice",
     "HasMatrixTranspose",
     "HasNDim",
     "HasShape",
@@ -16,6 +17,7 @@
 from ._array import (
     Array,
     HasArrayNamespace,
+    HasDevice,
     HasDType,
     HasMatrixTranspose,
     HasNDim,

src/array_api_typing/_array.py

@@ -16,6 +16,7 @@
 
 NamespaceT_co = TypeVar("NamespaceT_co", covariant=True, default=ModuleType)
 DTypeT_co = TypeVar("DTypeT_co", covariant=True)
+DeviceT_co = TypeVar("DeviceT_co", covariant=True, default=object)
 
 
 class HasArrayNamespace(Protocol[NamespaceT_co]):
@@ -74,11 +75,11 @@ def dtype(self, /) -> DTypeT_co:
         ...
 
 
-class HasDevice(Protocol):
+class HasDevice(Protocol[DeviceT_co]):
     """Protocol for array classes that have a device attribute."""
 
     @property
-    def device(self) -> object:  # TODO: more specific type
+    def device(self) -> DeviceT_co:
         """Hardware device the array data resides on."""
         ...
 
@@ -191,7 +192,7 @@ def T(self) -> Self:  # noqa: N802
 class Array(
     # ------ Attributes -------
     HasDType[DTypeT_co],
-    HasDevice,
+    HasDevice[DeviceT_co],
     HasMatrixTranspose,
     HasNDim,
     HasShape,
@@ -200,14 +201,18 @@ class Array(
     # ------- Methods ---------
     HasArrayNamespace[NamespaceT_co],
     # -------------------------
-    Protocol[DTypeT_co, NamespaceT_co],
+    Protocol[DTypeT_co, DeviceT_co, NamespaceT_co],
 ):
     """Array API specification for array object attributes and methods.
 
-    The type is: ``Array[+DTypeT, +NamespaceT = ModuleType] = Array[DTypeT,
-    NamespaceT]`` where:
+    The type is: ``Array[+DTypeT, +DeviceT = object, +NamespaceT = ModuleType] =
+    Array[DTypeT, DeviceT, NamespaceT]`` where:
 
     - `DTypeT` is the data type of the array elements.
+    - `DeviceT` is the type of the device attribute. It defaults to `object` to
+      enable skipping device specification. Array objects supporting device
+      management can specify a more specific type if they use types (as opposed
+      to object instances) to distinguish between different devices.
     - `NamespaceT` is the type of the array namespace. It defaults to
       `ModuleType`, which is the most common form of array namespace (e.g.,
       `numpy`, `cupy`, etc.). However, it can be any type, e.g. a

tests/integration/test_numpy1p0.pyi

@@ -43,7 +43,7 @@ _: xpt.HasDType[dtype[Any]] = nparr_f32
 # `xpt.Array`
 
 # Check NamespaceT_co assignment
-a_ns: xpt.Array[Any, ModuleType] = nparr
+a_ns: xpt.Array[Any, Any, ModuleType] = nparr
 
 # Check DTypeT_co assignment
 # Note that `np.array_api` uses dtype objects, not dtype classes, so we can't

tests/integration/test_numpy2p0.pyi

@@ -49,7 +49,7 @@ _: xpt.HasDType[np.dtype[B]] = nparr_b
 # `xpt.Array`
 
 # Check NamespaceT_co assignment
-a_ns: xpt.Array[Any, ModuleType] = nparr
+a_ns: xpt.Array[Any, Any, ModuleType] = nparr
 
 # Check DTypeT_co assignment
 _: xpt.Array[Any] = nparr
@@ -62,6 +62,13 @@ assert_type(x_f32.dtype, np.dtype[F32])
 assert_type(x_i32.dtype, np.dtype[I32])
 assert_type(x_b.dtype, np.dtype[B])
 
+# Check DeviceT_co assignment
+x_gooddevice: xpt.Array[Any, object, Any] = nparr
+assert_type(x_gooddevice.device, object)
+
+x_baddevice: xpt.Array[Any, int, Any] = nparr  # type: ignore[assignment]
+_: int = x_baddevice.device
+
 # Check Attribute `.device`
 assert_type(x_f32.device, object)
 assert_type(x_i32.device, object)