🚚 move HasArrayNamespace

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

Author: nstarman

.github/workflows/ci.yml

@@ -88,11 +88,14 @@ jobs:
           python-version: "3.11"
           activate-environment: true
 
-      - name: get major numpy version
-        id: numpy-major
+      - name: get major.minor numpy version
+        id: numpy-version
         run: |
-          version=$(echo ${{ matrix.numpy-version }} | cut -c 1)
-          echo "::set-output name=version::$version"
+          version="${{ matrix.numpy-version }}"
+          major=$(echo "$version" | cut -d. -f1)
+          echo "major=$major" >> $GITHUB_OUTPUT
+          minor=$(echo "$version" | cut -d. -f2)
+          echo "minor=$minor" >> $GITHUB_OUTPUT
 
       - name: install deps
         run: |
@@ -101,10 +104,21 @@ jobs:
 
       # NOTE: `uv run --with=...` will be ignored by mypy (and `--isolated` does not help)
       - name: mypy
-        run: >
-          uv run --no-sync --active
-          mypy --tb --no-incremental --cache-dir=/dev/null
-          tests/integration/test_numpy${{ steps.numpy-major.outputs.version }}.pyi
+        run: |
+          major="${{ steps.numpy-version.outputs.major }}"
+          minor="${{ steps.numpy-version.outputs.minor }}"
+
+          if [ "$major" -eq 1 ]; then
+            file="test_numpy1.pyi"
+          elif [ "$major" -eq 2 ] && [ "$minor" -lt 2 ]; then
+            file="test_numpy2.0.pyi"
+          else
+            file="test_numpy2.2.pyi"
+          fi
+
+          uv run --no-sync --active \
+            mypy --tb --no-incremental --cache-dir=/dev/null \
+            tests/integration/$file
 
       # TODO: (based)pyright
 

src/array_api_typing/__init__.py

@@ -6,5 +6,5 @@
     "__version_tuple__",
 )
 
-from ._namespace import HasArrayNamespace
+from ._array import HasArrayNamespace
 from ._version import version as __version__, version_tuple as __version_tuple__

src/array_api_typing/_namespace.py renamed to src/array_api_typing/_array.py

@@ -4,12 +4,17 @@
 from typing import Literal, Protocol
 from typing_extensions import TypeVar
 
-T_co = TypeVar("T_co", covariant=True, default=ModuleType)
+NamespaceT_co = TypeVar("NamespaceT_co", covariant=True, default=ModuleType)
 
 
-class HasArrayNamespace(Protocol[T_co]):
+class HasArrayNamespace(Protocol[NamespaceT_co]):
     """Protocol for classes that have an `__array_namespace__` method.
 
+    This `Protocol` is intended for use in static typing to ensure that an
+    object has an `__array_namespace__` method that returns a namespace for
+    array operations. This `Protocol` should not be used at runtime, for type
+    checking or as a base class.
+
     Example:
     >>> import array_api_typing as xpt
     >>>
@@ -27,4 +32,4 @@ class HasArrayNamespace(Protocol[T_co]):
 
     def __array_namespace__(
         self, /, *, api_version: Literal["2021.12"] | None = None
-    ) -> T_co: ...
+    ) -> NamespaceT_co: ...

tests/integration/test_numpy1.pyi

@@ -1,12 +1,34 @@
-from typing import Any
+# mypy: disable-error-code="no-redef"
 
-# requires numpy < 2
-import numpy.array_api as np
+from types import ModuleType
+from typing import TypeAlias
+
+import numpy.array_api as np  # type: ignore[import-not-found, unused-ignore]
 
 import array_api_typing as xpt
 
-###
-# Ensure that `np.ndarray` instances are assignable to `xpt.HasArrayNamespace`.
+# DType aliases
+F32: TypeAlias = np.float32
+I32: TypeAlias = np.int32
+
+# Define NDArrays against which we can test the protocols
+nparr = np.eye(2)
+nparr_i32 = np.array([1], dtype=I32)
+nparr_f32 = np.array([1.0], dtype=F32)
+nparr_b = np.array([True], dtype=np.bool_)
+
+# =========================================================
+# `xpt.HasArrayNamespace`
+
+_: xpt.HasArrayNamespace[ModuleType] = nparr
+_: xpt.HasArrayNamespace[ModuleType] = nparr_i32
+_: xpt.HasArrayNamespace[ModuleType] = nparr_f32
+_: xpt.HasArrayNamespace[ModuleType] = nparr_b
+
+# Check `__array_namespace__` method
+a_ns: xpt.HasArrayNamespace[ModuleType] = nparr
+ns: ModuleType = a_ns.__array_namespace__()
 
-arr = np.eye(2)
-arr_namespace: xpt.HasArrayNamespace[Any] = arr
+# Incorrect values are caught when using `__array_namespace__` and
+# backpropagated to the type of `a_ns`
+_: xpt.HasArrayNamespace[dict[str, int]] = nparr  # not caught

tests/integration/test_numpy2.0.pyi

@@ -0,0 +1,36 @@
+# mypy: disable-error-code="no-redef"
+
+from types import ModuleType
+from typing import Any, TypeAlias
+
+import numpy as np
+import numpy.typing as npt
+
+import array_api_typing as xpt
+
+# DType aliases
+F32: TypeAlias = np.float32
+I32: TypeAlias = np.int32
+
+# Define NDArrays against which we can test the protocols
+nparr: npt.NDArray[Any]
+nparr_i32: npt.NDArray[I32] = np.array([1], dtype=I32)
+nparr_f32: npt.NDArray[F32] = np.array([1.0], dtype=F32)
+nparr_b: npt.NDArray[np.bool_] = np.array([True], dtype=np.bool_)
+
+# =========================================================
+# `xpt.HasArrayNamespace`
+
+# Check assignment
+_: xpt.HasArrayNamespace[ModuleType] = nparr
+_: xpt.HasArrayNamespace[ModuleType] = nparr_i32
+_: xpt.HasArrayNamespace[ModuleType] = nparr_f32
+_: xpt.HasArrayNamespace[ModuleType] = nparr_b
+
+# Check `__array_namespace__` method
+a_ns: xpt.HasArrayNamespace[ModuleType] = nparr
+ns: ModuleType = a_ns.__array_namespace__()
+
+# Incorrect values are caught when using `__array_namespace__` and
+# backpropagated to the type of `a_ns`
+_: xpt.HasArrayNamespace[dict[str, int]] = nparr  # not caught

tests/integration/test_numpy2.2.pyi

@@ -0,0 +1,11 @@
+# mypy: disable-error-code="no-redef"
+
+from test_numpy2 import nparr
+
+import array_api_typing as xpt
+
+# Incorrect values are caught when using `__array_namespace__` and
+# backpropagated to the type of `a_ns`
+_: xpt.HasArrayNamespace[dict[str, int]] = nparr  # not caught
+a_badns: xpt.HasArrayNamespace[dict[str, int]] = nparr  # type: ignore[assignment]
+a_badns.__array_namespace__()  # triggers error above