Add protocols for array object methods

- DLPack export via `HasDLPack`
- Indexing via `HasGetItem`
- Indexed assignment via `HasSetItem`
- Device transfer via `HasToDevice`

Also composes these protocols into `Array`, exports them from the
top-level package, and adds NumPy integration assignments for the new
protocol surfaces.

Author: NeilGirdhar

src/array_api_typing/__init__.py

@@ -4,25 +4,9 @@
     "Array",
     "HasArrayNamespace",
     "HasDType",
-    "HasDevice",
-    "HasMatrixTranspose",
-    "HasNDim",
-    "HasShape",
-    "HasSize",
-    "HasTranspose",
     "__version__",
     "__version_tuple__",
 )
 
-from ._array import (
-    Array,
-    HasArrayNamespace,
-    HasDevice,
-    HasDType,
-    HasMatrixTranspose,
-    HasNDim,
-    HasShape,
-    HasSize,
-    HasTranspose,
-)
+from ._array import Array, HasArrayNamespace, HasDType
 from ._version import version as __version__, version_tuple as __version_tuple__

src/array_api_typing/_array.py

@@ -1,13 +1,6 @@
 __all__ = (
     "Array",
     "HasArrayNamespace",
-    "HasDType",
-    "HasDevice",
-    "HasMatrixTranspose",
-    "HasNDim",
-    "HasShape",
-    "HasSize",
-    "HasTranspose",
 )
 
 from types import ModuleType
@@ -17,6 +10,8 @@
 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)
+KeyT_contra = TypeVar("KeyT_contra", contravariant=True, default=object)
+ValueT_contra = TypeVar("ValueT_contra", contravariant=True, default=object)
 
 
 class HasArrayNamespace(Protocol[NamespaceT_co]):
@@ -66,6 +61,26 @@ def __array_namespace__(
         ...
 
 
+class HasDLPack(Protocol):
+    """Protocol for array classes that support DLPack export."""
+
+    def __dlpack__(
+        self,
+        /,
+        *,
+        stream: object | None = None,
+        max_version: tuple[int, int] | None = None,
+        dl_device: tuple[int, int] | None = None,
+        copy: bool | None = None,
+    ) -> object:
+        """Export the array as a DLPack capsule."""
+        ...
+
+    def __dlpack_device__(self, /) -> tuple[int, int]:
+        """Return the DLPack device type and device ID."""
+        ...
+
+
 class HasDType(Protocol[DTypeT_co]):
     """Protocol for array classes that have a data type attribute."""
 
@@ -84,6 +99,14 @@ def device(self) -> DeviceT_co:
         ...
 
 
+class HasGetItem(Protocol[KeyT_contra]):
+    """Protocol for array classes that support indexing."""
+
+    def __getitem__(self, key: KeyT_contra, /) -> Self:
+        """Return ``self[key]``."""
+        ...
+
+
 class HasMatrixTranspose(Protocol):
     """Protocol for array classes that have a matrix transpose attribute."""
 
@@ -140,6 +163,14 @@ def shape(self) -> tuple[int | None, ...]:
         ...
 
 
+class HasSetItem(Protocol[KeyT_contra, ValueT_contra]):
+    """Protocol for mutable array classes that support indexed assignment."""
+
+    def __setitem__(self, key: KeyT_contra, value: ValueT_contra, /) -> None:
+        """Set ``self[key]`` to ``value``."""
+        ...
+
+
 class HasSize(Protocol):
     """Protocol for array classes that have a size attribute."""
 
@@ -159,6 +190,14 @@ def size(self) -> int | None:
         ...
 
 
+class HasToDevice(Protocol):
+    """Protocol for array classes that support device transfer."""
+
+    def to_device(self, device: object, /, *, stream: object | None = None) -> Self:
+        """Copy the array to the specified device."""
+        ...
+
+
 class HasTranspose(Protocol):
     """Protocol for array classes that support the transpose operation."""
 
@@ -191,16 +230,26 @@ def T(self) -> Self:  # noqa: N802
 
 class Array(
     # ------ Attributes -------
+    HasDevice[DeviceT_co],
     HasDType[DTypeT_co],
+    HasMatrixTranspose,
+    HasNDim,
+    HasShape,
+    HasSize,
+    HasTranspose,
     # ------- Methods ---------
     HasArrayNamespace[NamespaceT_co],
+    HasDLPack,
+    HasGetItem,
+    HasSetItem,
+    HasToDevice,
     # -------------------------
-    Protocol[DTypeT_co, NamespaceT_co],
+    Protocol[DTypeT_co, NamespaceT_co, DeviceT_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, +NamespaceT = ModuleType, +DeviceT = object] =
+    Array[DTypeT, NamespaceT, DeviceT]`` where:
 
     - `DTypeT` is the data type of the array elements.
     - `NamespaceT` is the type of the array namespace. It defaults to
@@ -209,6 +258,7 @@ class Array(
       `types.SimpleNamespace`, to allow for wrapper libraries to
       semi-dynamically define their own array namespaces based on the wrapped
       array type.
+    - `DeviceT` is the type of the hardware device the array data resides on.
 
     This type is intended for use in static typing to ensure that an object has
     the attributes and methods defined in the array API specification. It should

tests/integration/test_numpy2p0.pyi

@@ -36,6 +36,18 @@ ns: ModuleType = a_ns.__array_namespace__()
 # backpropagated to the type of `a_ns`
 _: xpt.HasArrayNamespace[dict[str, int]] = nparr  # not caught
 
+# =========================================================
+# `xpt.HasDLPack`
+
+_: xpt.HasDLPack = nparr
+_: xpt.HasDLPack = nparr_i32
+_: xpt.HasDLPack = nparr_f32
+_: xpt.HasDLPack = nparr_b
+
+a_dlpack: xpt.HasDLPack = nparr
+assert_type(a_dlpack.__dlpack__(), object)
+assert_type(a_dlpack.__dlpack_device__(), tuple[int, int])
+
 # =========================================================
 # `xpt.HasDType`
 
@@ -53,6 +65,14 @@ _: xpt.HasDevice = nparr_i32
 _: xpt.HasDevice = nparr_f32
 _: xpt.HasDevice = nparr_b
 
+# =========================================================
+# `xpt.HasGetItem`
+
+_: xpt.HasGetItem = nparr
+_: xpt.HasGetItem = nparr_i32
+_: xpt.HasGetItem = nparr_f32
+_: xpt.HasGetItem = nparr_b
+
 # =========================================================
 # `xpt.HasMatrixTranspose`
 
@@ -77,6 +97,14 @@ _: xpt.HasShape = nparr_i32
 _: xpt.HasShape = nparr_f32
 _: xpt.HasShape = nparr_b
 
+# =========================================================
+# `xpt.HasSetItem`
+
+_: xpt.HasSetItem = nparr
+_: xpt.HasSetItem = nparr_i32
+_: xpt.HasSetItem = nparr_f32
+_: xpt.HasSetItem = nparr_b
+
 # =========================================================
 # `xpt.HasSize`
 
@@ -85,6 +113,14 @@ _: xpt.HasSize = nparr_i32
 _: xpt.HasSize = nparr_f32
 _: xpt.HasSize = nparr_b
 
+# =========================================================
+# `xpt.HasToDevice`
+
+_: xpt.HasToDevice = nparr
+_: xpt.HasToDevice = nparr_i32
+_: xpt.HasToDevice = nparr_f32
+_: xpt.HasToDevice = nparr_b
+
 # =========================================================
 # `xpt.HasTranspose`