✨HasDType

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

Author: nstarman

src/array_api_typing/__init__.py

@@ -3,9 +3,10 @@
 __all__ = (
     "Array",
     "HasArrayNamespace",
+    "HasDType",
     "__version__",
     "__version_tuple__",
 )
 
-from ._array import Array, HasArrayNamespace
+from ._array import Array, HasArrayNamespace, HasDType
 from ._version import version as __version__, version_tuple as __version_tuple__

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]):
@@ -57,8 +58,32 @@ def __array_namespace__(
         ...
 
 
+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_numpy1p0.pyi

@@ -1,9 +1,10 @@
 # mypy: disable-error-code="no-redef"
 
 from types import ModuleType
-from typing import TypeAlias
+from typing import Any
 
 import numpy.array_api as np  # type: ignore[import-not-found, unused-ignore]
+from numpy import dtype
 
 import array_api_typing as xpt
 
@@ -28,8 +29,25 @@ ns: ModuleType = a_ns.__array_namespace__()
 # backpropagated to the type of `a_ns`
 _: xpt.HasArrayNamespace[dict[str, int]] = nparr  # not caught
 
+# =========================================================
+# `xpt.HasDType`
+
+# Note that `np.array_api` uses dtype objects, not dtype classes, so we can't
+# type annotate specific dtypes like `np.float32` or `np.int32`.
+
+_: xpt.HasDType[dtype[Any]] = nparr
+_: xpt.HasDType[dtype[Any]] = nparr_i32
+_: xpt.HasDType[dtype[Any]] = nparr_f32
+
 # =========================================================
 # `xpt.Array`
 
 # Check NamespaceT_co assignment
-a_ns: xpt.Array[ModuleType] = nparr
+a_ns: xpt.Array[Any, ModuleType] = nparr
+
+# Check DTypeT_co assignment
+# Note that `np.array_api` uses dtype objects, not dtype classes, so we can't
+# type annotate specific dtypes like `np.float32` or `np.int32`.
+_: xpt.Array[dtype[Any]] = nparr
+_: xpt.Array[dtype[Any]] = nparr_i32
+_: xpt.Array[dtype[Any]] = nparr_f32

tests/integration/test_numpy2p0.pyi

@@ -35,8 +35,23 @@ ns: ModuleType = a_ns.__array_namespace__()
 # backpropagated to the type of `a_ns`
 _: xpt.HasArrayNamespace[dict[str, int]] = nparr  # not caught
 
+# =========================================================
+# `xpt.HasDType`
+
+# Check DTypeT_co assignment
+_: xpt.HasDType[Any] = nparr
+_: xpt.HasDType[np.dtype[I32]] = nparr_i32
+_: xpt.HasDType[np.dtype[F32]] = nparr_f32
+_: xpt.HasDType[np.dtype[np.bool_]] = nparr_b
+
 # =========================================================
 # `xpt.Array`
 
 # Check NamespaceT_co assignment
-a_ns: xpt.Array[ModuleType] = nparr
+a_ns: xpt.Array[Any, 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