✨ HasDType

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

Author: nstarman

pyproject.toml

@@ -130,6 +130,8 @@ ignore = [
   "ISC001", # Conflicts with formatter
   "PLW1641", # Object does not implement `__hash__` method
   "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

@@ -7,7 +7,7 @@
 
 from pathlib import Path
 from types import ModuleType
-from typing import Literal, Never, Protocol, TypeAlias
+from typing import Any, Literal, Never, Protocol, TypeAlias
 from typing_extensions import TypeVar
 
 import optype as op
@@ -27,6 +27,7 @@
 NS_co = TypeVar("NS_co", covariant=True, default=ModuleType)
 Other_contra = TypeVar("Other_contra", contravariant=True, default=Never)
 R_co = TypeVar("R_co", covariant=True, default=Never)
+DType_co = TypeVar("DType_co", covariant=True)
 
 
 class HasArrayNamespace(Protocol[NS_co]):
@@ -52,8 +53,20 @@ def __array_namespace__(
     ) -> NS_co: ...
 
 
+class HasDType(Protocol[DType_co]):
+    """Protocol for array classes that have a data type attribute."""
+
+    @property
+    def dtype(self) -> DType_co:
+        """Data type of the array elements."""
+        ...
+
+
 @docstring_setter(**_array_docstrings)
 class Array(
+    # ------ Attributes -------
+    HasDType[DType_co],
+    # ------ Methods -------
     HasArrayNamespace[NS_co],
     op.CanPosSelf,
     op.CanNegSelf,
@@ -64,22 +77,23 @@ class Array(
     op.CanFloordivSame[Other_contra, R_co],
     op.CanModSame[Other_contra, R_co],
     op.CanPowSame[Other_contra, R_co],
-    Protocol[Other_contra, R_co, NS_co],
+    Protocol[DType_co, Other_contra, R_co, NS_co],
 ):
     """Array API specification for array object attributes and methods.
 
-    The type is: ``Array[-Other = Never, +R = Never, +NS = ModuleType] =
-    Array[Self | Other, Self | R, NS]`` where:
+    The type is: ``Array[+DType, -Other = Never, +R = Never, +NS = ModuleType] =
+    Array[+DType, Self | Other, Self | R, NS]`` where:
 
+    - `DType` is the data type of the array elements.
     - `Other` is the type of objects that can be used with the array (e.g., for
       binary operations). For example, with numeric arrays, it is common to be
       able to add `float` and `int` objects to the array, not just other arrays
       of the same dtype. This would be annotated as `Other = float`.  When not
       specified, `Other` only allows `Self` objects.
     - `R` is the return type of the array operations. For example, the return
       type of the division between integer arrays can be a float array. This
-      would be annotated as `R = float`. When not specified, `R` only allows
-      `Self` objects.
+      would be annotated as `R = Array[float]`. When not specified, `R` only
+      allows `Self` objects.
     - `NS` 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
@@ -89,20 +103,47 @@ class Array(
     """
 
 
-BoolArray: TypeAlias = Array[bool, Array[float, Never, NS_co], NS_co]
-"""Array API specification for boolean array object attributes and methods.
+# TODO: are there ways to tighten the dtype in both Arrays?
+BoolArray: TypeAlias = Array[DType_co, bool, Array[Any, float, Never, NS_co], NS_co]
+"""Array API specification for arrays that work with boolean values.
 
-Specifically, this type alias fills the `Other_contra` type variable with
-`bool`, allowing for `bool` objects to be added, subtracted, multiplied, etc. to
-the array object.
+The type is: ``BoolArray[+DType, +NS = ModuleType] = Array[+DType, Self | bool,
+Self | Array[Any, float, Self, NS], Self, NS]`` where:
 
-"""
+- `DType` is the data type of the array elements.
+- The second type variable -- `Other` -- is filled with `bool`, allowing for
+  `bool` objects to be added, subtracted, multiplied, etc. to the array object.
+- The third type variable -- `R` -- is filled with `Array[Any, float, Self,
+  NS]`, which is the return type of the array operations. For example, the
+  return type of the division between boolean arrays can be a float array.
+- `NS` 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.
 
-NumericArray: TypeAlias = Array[float | int, NS_co]
-"""Array API specification for numeric array object attributes and methods.
+"""
 
-Specifically, this type alias fills the `Other_contra` type variable with `float
-| int`, allowing for `float | int` objects to be added, subtracted, multiplied,
-etc. to the array object.
+# TODO: are there ways to tighten the dtype in both Arrays?
+NumericArray: TypeAlias = Array[
+    DType_co, float | int, Array[Any, float | int, Never, NS_co], NS_co
+]
+"""Array API specification for arrays that work with numeric values.
+
+the type is: ``NumericArray[+DType, +NS = ModuleType] = Array[+DType, Self |
+float | int, Self | Array[Any, float | int, Self, NS], NS]`` where:
+
+- `DType` is the data type of the array elements.
+- The second type variable -- `Other` -- is filled with `float | int`, allowing
+  for `float | int` objects to be added, subtracted, multiplied, etc. to the
+  array object.
+- The third type variable -- `R` -- is filled with `Array[Any, float | int,
+  Self, NS]`, which is the return type of the array operations. For example, the
+  return type of the division between integer arrays can be a float array.
+- `NS` 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

@@ -42,4 +42,4 @@ arr_i: xpt.Array[int | float, xpt.Array[float | int]] = nparr_i32
 # =========================================================
 # Check np.ndarray against BoolArray and NumericArray type aliases
 
-numericarray: NumericArray = nparr_f32
+numericarray: NumericArray[Any] = nparr_f32

tests/integration/test_numpy2.pyi

@@ -1,15 +1,23 @@
-from typing import Any, Never, TypeAlias
+from typing import Any, TypeAlias
 
 import numpy as np
 import numpy.typing as npt
 
 import array_api_typing as xpt
-from array_api_typing._array import BoolArray, NumericArray
+from array_api_typing._array import BoolArray, HasDType, NumericArray
 
 F: TypeAlias = np.floating[Any]
+DtF: TypeAlias = np.dtype[F]
 F32: TypeAlias = np.float32
+DtF32: TypeAlias = np.dtype[F32]
+
 I: TypeAlias = np.integer[Any]
+DtI: TypeAlias = np.dtype[I]
 I32: TypeAlias = np.int32
+DtI32: TypeAlias = np.dtype[I32]
+
+DtB: TypeAlias = np.dtype[np.bool_]
+DtN: TypeAlias = np.dtype[np.number]
 
 # Define NDArrays against which we can test the protocols
 nparr: npt.NDArray[Any]
@@ -24,28 +32,35 @@ arr_ns: xpt.HasArrayNamespace[Any] = nparr
 arr_ns_i32: xpt.HasArrayNamespace[Any] = nparr_i32
 arr_ns_f32: xpt.HasArrayNamespace[Any] = nparr_f32
 
+# =========================================================
+# Ensure that `np.ndarray` instances are assignable to `xpt.HasDType`.
+
+arr_dtype: HasDType[Any] = nparr
+arr_dtype_i32: HasDType[DtI32] = nparr_i32
+arr_dtype_f32: HasDType[DtF32] = nparr_f32
+
 # =========================================================
 # Ensure that `np.ndarray` instances are assignable to `xpt.Array`.
 
 # Generic Array type
-arr_array: xpt.Array[Never] = nparr
+arr_array: xpt.Array[Any] = nparr
 
 # Float Array types
-arr_float: xpt.Array[float] = nparr_f32
-arr_f: xpt.Array[F] = nparr_f32
-arr_f32: xpt.Array[F32] = nparr_f32
+arr_float: xpt.Array[DtF, float] = nparr_f32
+arr_f: xpt.Array[DtF, F] = nparr_f32
+arr_f32: xpt.Array[DtF32, F32] = nparr_f32
 
 # Integer Array types
-arr_int: xpt.Array[int, xpt.Array[float | int]] = nparr_i32
-arr_i: xpt.Array[I, xpt.Array[float | int]] = nparr_i32
-arr_i32: xpt.Array[I32, xpt.Array[F32 | I32]] = nparr_i32
+arr_int: xpt.Array[DtI32, int, xpt.Array[DtN, float | int]] = nparr_i32
+arr_i: xpt.Array[DtI32, I, xpt.Array[DtN, float | int]] = nparr_i32
+arr_i32: xpt.Array[DtI32, I32, xpt.Array[DtN, F32 | I32]] = nparr_i32
 
 # Boolean Array types
-arr_bool: xpt.Array[bool, xpt.Array[float | int | bool]] = nparr_b
-arr_b: xpt.Array[np.bool_, xpt.Array[F | I | np.bool_]] = nparr_b
+arr_bool: xpt.Array[DtB, bool, xpt.Array[DtN | DtB, float | int | bool]] = nparr_b
+arr_b: xpt.Array[DtB, np.bool_, xpt.Array[DtN | DtB, F | I | np.bool_]] = nparr_b
 
 # =========================================================
 # Check np.ndarray against BoolArray and NumericArray type aliases
 
-boolarray: BoolArray = nparr_b
-numericarray: NumericArray = nparr_f32
+boolarray: BoolArray[DtB] = nparr_b  # TODO: the Other type isn't correct
+numericarray: NumericArray[DtN] = nparr_f32