feat: HasDType

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

Author: nstarman

src/array_api_typing/_array.py

@@ -25,8 +25,9 @@
     _array_docstrings = tomllib.load(f)["docstrings"]
 
 NS_co = TypeVar("NS_co", covariant=True, default=ModuleType)
-T_contra = TypeVar("T_contra", contravariant=True)
+Other_contra = TypeVar("Other_contra", contravariant=True)
 R_co = TypeVar("R_co", covariant=True, default=Never)
+DType_co = TypeVar("DType_co", covariant=True)
 
 
 class HasArrayNamespace(Protocol[NS_co]):
@@ -52,27 +53,60 @@ 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,
-    op.CanAddSame[T_contra, R_co],
-    op.CanSubSame[T_contra, R_co],
-    op.CanMulSame[T_contra, R_co],
-    op.CanTruedivSame[T_contra, R_co],
-    op.CanFloordivSame[T_contra, R_co],
-    op.CanModSame[T_contra, R_co],
-    op.CanPowSame[T_contra, R_co],
-    Protocol[T_contra, R_co, NS_co],
+    op.CanAddSame[Other_contra, R_co],
+    op.CanSubSame[Other_contra, R_co],
+    op.CanMulSame[Other_contra, R_co],
+    op.CanTruedivSame[Other_contra, R_co],
+    op.CanFloordivSame[Other_contra, R_co],
+    op.CanModSame[Other_contra, R_co],
+    op.CanPowSame[Other_contra, R_co],
+    Protocol[DType_co, Other_contra, R_co, NS_co],
 ):
-    """Array API specification for array object attributes and methods."""
+    """Array API specification for array object attributes and methods.
+
+    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.
+    - `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.
+
+    """
 
 
 BoolArray: TypeAlias = Array[bool, Array[float, Never, NS_co], NS_co]
 """Array API specification for boolean array object attributes and methods.
 
-Specifically, this type alias fills the `T_contra` type variable with
+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.
 
@@ -81,7 +115,7 @@ class Array(
 NumericArray: TypeAlias = Array[float | int, NS_co]
 """Array API specification for numeric array object attributes and methods.
 
-Specifically, this type alias fills the `T_contra` type variable with `float
+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.