✨ HasDType

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

Author: nstarman

pyproject.toml

@@ -124,6 +124,9 @@ ignore = [
   "FBT", # flake8-boolean-trap
   "FIX", # flake8-fixme
   "ISC001", # Conflicts with formatter
+  "PYI041", # Use `float` instead of `int | float`
+  "TD002", # Missing author in TODO
+  "TD003", # Missing issue link for this TODO
 ]
 
 [tool.ruff.lint.pylint]

src/array_api_typing/_array.py

@@ -8,6 +8,7 @@
 from typing_extensions import TypeVar
 
 NamespaceT_co = TypeVar("NamespaceT_co", covariant=True, default=ModuleType)
+DTypeT_co = TypeVar("DTypeT_co", covariant=True)
 
 
 class HasArrayNamespace(Protocol[NamespaceT_co]):
@@ -38,8 +39,32 @@ def __array_namespace__(
     ) -> NamespaceT_co: ...
 
 
+class HasDType(Protocol[DTypeT_co]):
+    """Protocol for array classes that have a data type attribute."""
+
+    @property
+    def dtype(self) -> DTypeT_co:
+        """Data type of the array elements."""
+        ...
+
+
 class Array(
-    HasArrayNamespace[NamespaceT_co],
-    Protocol[NamespaceT_co],
+    # ------ Attributes -------
+    HasDType[DTypeT_co],
+    # -------------------------
+    Protocol[DTypeT_co, NamespaceT_co],
 ):
-    """Array API specification for array object attributes and methods."""
+    """Array API specification for array object attributes and methods.
+
+    The type is: ``Array[+DTypeT, +NamespaceT = ModuleType] = Array[DTypeT,
+    NamespaceT]`` where:
+
+    - `DTypeT` is the data type of the array elements.
+    - `NamespaceT` is the type of the array namespace. It defaults to
+      `ModuleType`, which is the most common form of array namespace (e.g.,
+      `numpy`, `cupy`, etc.). However, it can be any type, e.g. a
+      `types.SimpleNamespace`, to allow for wrapper libraries to
+      semi-dynamically define their own array namespaces based on the wrapped
+      array type.
+
+    """

tests/integration/test_numpy1.pyi

@@ -1,7 +1,7 @@
 # mypy: disable-error-code="no-redef"
 
 from types import ModuleType
-from typing import TypeAlias
+from typing import Any, TypeAlias
 
 import numpy.array_api as np  # type: ignore[import-not-found, unused-ignore]
 
@@ -38,3 +38,9 @@ _: xpt.HasArrayNamespace[dict[str, int]] = nparr  # not caught
 
 # Check NamespaceT_co assignment
 a_ns: xpt.Array[ModuleType] = nparr
+
+# Check DTypeT_co assignment
+_: xpt.Array[Any] = nparr
+_: xpt.Array[np.dtype[I32]] = nparr_i32
+_: xpt.Array[np.dtype[F32]] = nparr_f32
+_: xpt.Array[np.dtype[np.bool_]] = nparr_b

tests/integration/test_numpy2.0.pyi

@@ -9,7 +9,9 @@ import numpy.typing as npt
 import array_api_typing as xpt
 
 # DType aliases
+F: TypeAlias = np.floating[Any]
 F32: TypeAlias = np.float32
+I: TypeAlias = np.integer[Any]
 I32: TypeAlias = np.int32
 
 # Define NDArrays against which we can test the protocols
@@ -40,3 +42,9 @@ _: xpt.HasArrayNamespace[dict[str, int]] = nparr  # not caught
 
 # Check NamespaceT_co assignment
 a_ns: xpt.Array[ModuleType] = nparr
+
+# Check DTypeT_co assignment
+_: xpt.Array[Any] = nparr
+_: xpt.Array[np.dtype[I32]] = nparr_i32
+_: xpt.Array[np.dtype[F32]] = nparr_f32
+_: xpt.Array[np.dtype[np.bool_]] = nparr_b