✨ HasDevice

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

Author: nstarman

pyproject.toml

@@ -126,6 +126,8 @@ ignore = [
   "FIX", # flake8-fixme
   "ISC001", # Conflicts with formatter
   "PYI041", # Use `float` instead of `int | float`
+  "TD002", # Missing author in TODO
+  "TD003", # Missing issue link for this TODO
 ]
 
 [tool.ruff.lint.pydocstyle]

src/array_api_typing/__init__.py

@@ -4,13 +4,15 @@
     "Array",
     "HasArrayNamespace",
     "HasDType",
+    "HasDevice",
     "__version__",
     "__version_tuple__",
 )
 
 from ._array import (
     Array,
     HasArrayNamespace,
+    HasDevice,
     HasDType,
 )
 from ._version import version as __version__, version_tuple as __version_tuple__

src/array_api_typing/_array.py

@@ -2,6 +2,7 @@
     "Array",
     "HasArrayNamespace",
     "HasDType",
+    "HasDevice",
 )
 
 from types import ModuleType
@@ -68,9 +69,19 @@ def dtype(self, /) -> DTypeT_co:
         ...
 
 
+class HasDevice(Protocol):
+    """Protocol for array classes that have a device attribute."""
+
+    @property
+    def device(self) -> object:  # TODO: more specific type
+        """Hardware device the array data resides on."""
+        ...
+
+
 class Array(
     # ------ Attributes -------
     HasDType[DTypeT_co],
+    HasDevice,
     # ------- Methods ---------
     HasArrayNamespace[NamespaceT_co],
     # -------------------------

tests/integration/test_numpy1p0.pyi

@@ -55,3 +55,7 @@ x_i32: xpt.Array[dtype[Any]] = nparr_i32
 # Check Attribute `.dtype`
 _: dtype[Any] = x_f32.dtype
 _: dtype[Any] = x_i32.dtype
+
+# Check Attribute `.device`
+_: object = x_f32.device
+_: object = x_i32.device

tests/integration/test_numpy2p0.pyi

@@ -61,3 +61,8 @@ x_b: xpt.Array[np.dtype[B]] = nparr_b
 _: np.dtype[F32] = x_f32.dtype
 _: np.dtype[I32] = x_i32.dtype
 _: np.dtype[B] = x_b.dtype
+
+# Check Attribute `.device`
+_: object = x_f32.device
+_: object = x_i32.device
+_: object = x_b.device