✨ 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,6 +69,15 @@ 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],

tests/integration/test_numpy1p0.pyi

@@ -39,6 +39,13 @@ _: xpt.HasDType[dtype[Any]] = nparr
 _: xpt.HasDType[dtype[Any]] = nparr_i32
 _: xpt.HasDType[dtype[Any]] = nparr_f32
 
+# =========================================================
+# `xpt.HasDevice`
+
+_: xpt.HasDevice = nparr
+_: xpt.HasDevice = nparr_i32
+_: xpt.HasDevice = nparr_f32
+
 # =========================================================
 # `xpt.Array`
 

tests/integration/test_numpy2p0.pyi

@@ -45,6 +45,14 @@ _: xpt.HasDType[np.dtype[I32]] = nparr_i32
 _: xpt.HasDType[np.dtype[F32]] = nparr_f32
 _: xpt.HasDType[np.dtype[B]] = nparr_b
 
+# =========================================================
+# `xpt.HasDevice`
+
+_: xpt.HasDevice = nparr
+_: xpt.HasDevice = nparr_i32
+_: xpt.HasDevice = nparr_f32
+_: xpt.HasDevice = nparr_b
+
 # =========================================================
 # `xpt.Array`