✅: dask

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

Author: nstarman

.github/workflows/ci.yml

@@ -144,5 +144,47 @@ jobs:
 
       # TODO: (based)pyright
 
+  test_integration_dask:
+    name: integration tests (dask)
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        dask-version: ["2025.7.0"]
+
+    steps:
+      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
+
+      - uses: astral-sh/setup-uv@d9e0f98d3fc6adb07d1e3d37f3043649ddad06a1 # v6.5.0
+        with:
+          python-version: "3.11"
+          activate-environment: true
+
+      # NOTE: `uv run --with=...` will be ignored by mypy (and `--isolated` does not help)
+      - name: mypy
+        run: |
+          uv sync --no-editable --group=mypy
+          uv pip install dask==${{ matrix.dask-version }}
+          uv run --no-sync --active \
+            mypy --tb --no-incremental --cache-dir=/dev/null \
+            tests/integration/test_dask.pyi
+
+      - name: basedmypy
+        run: |
+          uv sync --no-editable --group=basedmypy
+          uv pip install dask==${{ matrix.dask-version }}
+          uv run --no-sync --active \
+            mypy --tb --no-incremental --cache-dir=/dev/null \
+            tests/integration/test_dask.pyi
+
+      - name: pyright
+        run: |
+          uv sync --no-editable --group=pyright
+          uv pip install dask==${{ matrix.dask-version }}
+          uv run --no-sync --active \
+            tests/integration/test_dask.pyi
+
+      # TODO: (based)pyright
+
   # TODO: integration tests for array-api-strict
-  # TODO: integration tests for 3rd party libs such as cupy, pytorch, tensorflow, dask, etc.
+  # TODO: integration tests for 3rd party libs such as cupy, pytorch, tensorflow, etc.

pyproject.toml

@@ -68,6 +68,9 @@ pyright = [
 basedmypy = [
   "basedmypy>=2.10.0",
 ]
+test-dask = [
+  "dask>=2025.7.0",
+]
 
 [tool.hatch]
 version.source = "vcs"

tests/integration/test_dask.pyi

@@ -0,0 +1,54 @@
+# mypy: disable-error-code="no-redef, explicit-any"
+
+from types import ModuleType
+from typing import Any, TypeAlias
+
+import dask.array as da
+import numpy as np
+
+import array_api_typing as xpt
+
+# DType aliases
+F32: TypeAlias = np.float32
+I32: TypeAlias = np.int32
+
+# Define Dask Arrays against which we can test the protocols
+# (Dask's Array type is not parametrized by dtype at type level.)
+darr: da.Array
+darr_i32: da.Array
+darr_f32: da.Array
+darr_b: da.Array
+
+# =========================================================
+# `xpt.HasArrayNamespace`
+
+# Check assignment
+_001: xpt.HasArrayNamespace[ModuleType] = darr
+_002: xpt.HasArrayNamespace[ModuleType] = darr_i32
+_003: xpt.HasArrayNamespace[ModuleType] = darr_f32
+_004: xpt.HasArrayNamespace[ModuleType] = darr_b
+
+# Check `__array_namespace__` method
+a_ns: xpt.HasArrayNamespace[ModuleType] = darr
+ns: ModuleType = a_ns.__array_namespace__()
+
+# =========================================================
+# `xpt.HasDType`
+
+# Check DTypeT_co assignment
+_005: xpt.HasDType[Any] = darr
+_006: xpt.HasDType[np.dtype[I32]] = darr_i32
+_007: xpt.HasDType[np.dtype[F32]] = darr_f32
+_008: xpt.HasDType[np.dtype[np.bool_]] = darr_b
+
+# =========================================================
+# `xpt.Array`
+
+# Check NamespaceT_co assignment
+x_ns: xpt.Array[Any, ModuleType] = darr
+
+# Check DTypeT_co assignment
+_009: xpt.Array[Any] = darr
+_010: xpt.Array[np.dtype[I32]] = darr_i32
+_011: xpt.Array[np.dtype[F32]] = darr_f32
+_012: xpt.Array[np.dtype[np.bool_]] = darr_b