🚚 move HasArrayNamespace

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

Author: nstarman

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,36 @@
-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
+a_badns: xpt.HasArrayNamespace[dict[str, int]] = nparr  # type: ignore[assignment]
+a_badns.__array_namespace__()  # triggers error above

tests/integration/test_numpy2.pyi

@@ -1,11 +1,38 @@
-from typing import Any
+# 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
 
-###
-# 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: 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__()
 
-arr: npt.NDArray[Any]
-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
+a_badns: xpt.HasArrayNamespace[dict[str, int]] = nparr  # type: ignore[assignment]
+a_badns.__array_namespace__()  # triggers error above