Repository: jorenham/optype (language: python)

Issue:
Array-API interfaces
https://data-apis.org/array-api/latest/API_specification/index.html

## goal

The goal is to provide `Protocols` that can be used to type e.g. [`scipy.fft.fftfreq`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.fft.fftfreq.html#scipy.fft.fftfreq) (through its `xp` kwarg) and [`scipy.fft.fftshift`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.fft.fftshift.html) (via the passed array).

## the (big) problems

- The dtype spec only defines `__eq__`, which also exists in `object`. So it's *impossible* to (structurally) type a "dtype".
- In the same way, it's impossible to (structurally) type a "device".
- Shapes are defined as tuples of "something". So the best we can do is type the number of dimensions, and ignore the axis types.

This means that static type-checkers won't be able to see the difference between a boolean array on your CPU, and a complex-floating array on a quantum computer within the core of a neutron star in the Andromeda galaxy.

## so now what?

We can still figure out the following types:

- `Array` + the dimensionality (because `shape: tuple[?, ...]` kinda tells us the `ndim`, but not the actual shape)
- `Namespace`, which is a protocollable module-type (often named `xp`), that could give us 
    - `xp.array`: the specific `Array` type (but there's still no HKT after asking for it for over 6 years...), and 
    - `xp.(bool|((u?)int|float|complex)(8|16|32|64))` -i.e. the impossible-to-type `DType` types (🎩 🪄 🐰)
    - `xp.__array_namespace_info__() -> Info` ([docs](https://data-apis.org/array-api/latest/API_specification/inspection.html)), an object with some other "magical" methods:
        - `.default_device() -> DefaultDevice`  
        - `.devices() -> list[Device]`
        - `.dtypes(*, device: Device?, kind: Kind?)` - with `Kind` a literal string category, which allows us to "categorize" the e.g. floats, ints, and complexes (assuming the lib properly typed it), i.e. as type aliases.

## problem solved: just type the untypeable

Note that we can also obtain the `Info` type via `Array`, since `Array.__array_namespace__() -> Namespace[Info]` s.t. `Namespace[InfoT].__array_namespace_info__() -> InfoT` (ignoring the versions). 
So we can get all dtype types from either `Array` or `Namespace`, and we can either one from the other, and vice-versa (i.e. bidirectionally in both ways, and also the other way around).



Discussion / hints:

Fix patch:
diff --git a/.vscode/settings.json b/.vscode/settings.json
index 4ebdde06..95c09bd8 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,16 +1,14 @@
   {
-  "[python]": {
-    "editor.defaultFormatter": "charliermarsh.ruff",
-    "editor.formatOnPaste": true,
-    "editor.formatOnSave": true,
-    "editor.formatOnSaveMode": "modificationsIfAvailable",
-    "editor.formatOnType": true
-  },
-  "editor.rulers": [88],
-  "git.branchProtection": ["master"],
-  "mypy-type-checker.args": [
-    "--config-file=pyproject.toml"
-  ],
-  "python.testing.pytestEnabled": true,
-  "editor.defaultFormatter": "charliermarsh.ruff"
-}
+    "[python]": {
+      "editor.defaultFormatter": "charliermarsh.ruff",
+      "editor.formatOnPaste": true,
+      "editor.formatOnSave": true,
+      "editor.formatOnSaveMode": "modificationsIfAvailable",
+      "editor.formatOnType": true
+    },
+    "editor.rulers": [88],
+    "git.branchProtection": ["master"],
+    "mypy-type-checker.args": ["--config-file=pyproject.toml", "--ide"],
+    "python.testing.pytestEnabled": true,
+    "editor.defaultFormatter": "charliermarsh.ruff"
+  }
diff --git a/README.md b/README.md
index 6ccae3eb..6fefee20 100644
--- a/README.md
+++ b/README.md
@@ -225,6 +225,7 @@ The API of `optype` is flat; a single `import optype as opt` is all you need
     - [`Any*` type aliases](#any-type-aliases)
     - [`Empty*` type aliases](#empty-type-aliases)
     - [Literal types](#literal-types)
+- [`optype.dlpack`](#optypedlpack)
 - [`optype.numpy`](#optypenumpy)
     - [`Array`](#array)
     - [`UFunc`](#ufunc)
@@ -2095,6 +2096,75 @@ yield no elements.
     </tr>
 </table>
 
+### `optype.dlpack`
+
+A collection of low-level types for working [DLPack](DOC-DLPACK).
+
+#### Protocols
+
+<table>
+    <tr>
+        <th>type signature</th>
+        <th>bound method</th>
+    </tr>
+    <tr>
+<td>
+
+```plain
+CanDLPack[
+    +T = int,
+    +D: int = int,
+]
+```
+
+</td>
+<td>
+
+```python
+def __dlpack__(
+    *,
+    stream: int | None = ...,
+    max_version: tuple[int, int] | None = ...,
+    dl_device: tuple[T, D] | None = ...,
+    copy: bool | None = ...,
+) -> types.CapsuleType: ...
+```
+
+</td>
+    </tr>
+    <tr></tr>
+    <tr>
+<td>
+
+```plain
+CanDLPackDevice[
+    +T = int,
+    +D: int = int,
+]
+```
+
+</td>
+<td>
+
+```python
+def __dlpack_device__() -> tuple[T, D]: ...
+```
+
+</td>
+    </tr>
+</table>
+
+The `+` prefix indicates that the type parameter is *co*variant.
+
+#### Enums
+
+There are also two convenient
+[`IntEnum`](https://docs.python.org/3/library/enum.html#enum.IntEnum)s
+in `optype.dlpack`: `DLDeviceType` for the device types, and `DLDataTypeCode` for the
+internal type-codes of the `DLPack` data types.
+
+<table>
+
 ### `optype.numpy`
 
 Optype supports both NumPy 1 and 2.
diff --git a/optype/dlpack.py b/optype/dlpack.py
new file mode 100644
index 00000000..66eea390
--- /dev/null
+++ b/optype/dlpack.py
@@ -0,0 +1,97 @@
+"""
+Types and interfaces for DLPack, as used the Array API.
+https://github.com/dmlc/dlpack
+"""
+
+from __future__ import annotations
+
+import enum
+import sys
+from typing import Any, runtime_checkable
+
+
+if sys.version_info >= (3, 13):
+    from types import CapsuleType
+    from typing import Protocol, TypeVar
+else:
+    from typing_extensions import CapsuleType, Protocol, TypeVar
+
+
+__all__ = ["CanDLPack", "CanDLPackDevice"]
+
+
+class DLDeviceType(enum.IntEnum):
+    # GPU device
+    CPU = 1
+    # Cuda GPU device
+    CUDA = 2
+    # Pinned CUDA GPU memory vy `cudaMallocHost`
+    CPU_PINNED = 3
+    # OpenCL devices
+    OPENCL = 4
+    # Vulkan buffer for next generation graphics
+    VULKAN = 7
+    # Metal for Apple GPU
+    METAL = 8
+    # Verilog simulation buffer
+    VPI = 9
+    # ROCm GPU's for AMD GPU's
+    ROCM = 10
+    # CUDA managed/unified memory allocated by `cudaMallocManaged`
+    CUDA_MANAGED = 13
+    # Unified shared memory allocated on a oneAPI non-partititioned
+    # device. Call to oneAPI runtime is required to determine the device
+    # type, the USM allocation type and the sycl context it is bound to.
+    ONE_API = 14
+
+
+class DLDataTypeCode(enum.IntEnum):
+    # signed integer
+    INT = 0
+    # unsigned integer
+    UINT = 1
+    # IEEE floating point
+    FLOAT = 2
+    # Opaque handle type, reserved for testing purposes.
+    OPAQUE_HANDLE = 3
+    # bfloat16
+    BFLOAT = 4
+    # complex number (C/C++/Python layout: compact struct per complex number)
+    COMPLEX = 5
+    # boolean
+    BOOL = 6
+
+
+_TypeT_co = TypeVar(
+    "_TypeT_co",
+    covariant=True,
+    bound=enum.Enum | int,
+    default=enum.Enum | int,
+)
+_DeviceT_co = TypeVar("_DeviceT_co", covariant=True, bound=int, default=int)
+
+
+# NOTE: Because `__dlpack__` doesn't mutate the type, and the type parameters bind to
+# the *co*variant `tuple`, they should be *co*variant; NOT *contra*variant!
+# NOTE NOTE: This shows that PEP 695 claim, i.e. that variance can always be inferred,
+# is utter nonsense.
+
+
+@runtime_checkable
+class CanDLPack(Protocol[_TypeT_co, _DeviceT_co]):  # type: ignore[misc] # pyright: ignore[reportInvalidTypeVarUse]
+    def __dlpack__(  # type: ignore[no-any-explicit]
+        self,
+        /,
+        *,
+        stream: int | None = None,
+        max_version: tuple[int, int] | None = None,
+        dl_device: tuple[_TypeT_co, _DeviceT_co] | None = None,
+        # NOTE: This should be `bool | None`, but because of an incorrect annotation in
+        # `numpy.ndarray.__dlpack__`, this is not possible at the moment.
+        copy: Any | None = None,
+    ) -> CapsuleType: ...
+
+
+@runtime_checkable
+class CanDLPackDevice(Protocol[_TypeT_co, _DeviceT_co]):
+    def __dlpack_device__(self, /) -> tuple[_TypeT_co, _DeviceT_co]: ...
diff --git a/pyproject.toml b/pyproject.toml
index caacfc07..3c8e8912 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -295,6 +295,10 @@ force-exclude = true
             "__typing_unpacked_tuple_args__",
             "__typing_is_unpacked_typevartuple__",
 
+            # `dlpack`
+            "__dlpack__",
+            "__dlpack_device__",
+
             # `numpy` special methods
             "__array__",
             "__array_finalize__",


Test patch:
diff --git a/tests/test_dlpack.py b/tests/test_dlpack.py
new file mode 100644
index 00000000..7b27da61
--- /dev/null
+++ b/tests/test_dlpack.py
@@ -0,0 +1,18 @@
+import numpy as np
+
+from optype.dlpack import CanDLPack, CanDLPackDevice
+
+
+def test_ndarray_can_dlpack():
+    x: np.ndarray[tuple[()], np.dtype[np.float64]] = np.empty(0)
+    # NOTE: mypy doesn't understand covariance...
+    x_dl: CanDLPack = x  # type: ignore[assignment]
+
+    assert isinstance(x, CanDLPack)
+
+
+def test_ndarray_can_dlpack_device():
+    x: np.ndarray[tuple[()], np.dtype[np.float64]] = np.empty(0)
+    x_dl: CanDLPackDevice = x
+
+    assert isinstance(x, CanDLPack)