feat: HasX attributes

pyproject.toml

@@ -27,6 +27,9 @@ classifiers = [
 ]
 dependencies = [
   "typing-extensions>=4.14.1",
+  "optype>=0.9.3; python_version < '3.11'",
+  "optype>=0.12.2; python_version >= '3.11'",
+  "tomli>=1.2.0 ; python_full_version < '3.11'",
 ]
 
 [project.urls]
@@ -121,9 +124,14 @@ ignore = [
   "D107", # Missing docstring in __init__
   "D203", # 1 blank line required before class docstring
   "D213", # Multi-line docstring summary should start at the second line
+  "D401", # First line of docstring should be in imperative mood
   "FBT", # flake8-boolean-trap
   "FIX", # flake8-fixme
   "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]
@@ -137,10 +145,13 @@ allow-dunder-method-names = [
 ]
 
 [tool.ruff.lint.flake8-import-conventions]
-banned-from = ["array_api_typing"]
+banned-from = ["array_api_typing", "optype", "optype.numpy", "optype.numpy.compat"]
 
 [tool.ruff.lint.flake8-import-conventions.extend-aliases]
 array_api_typing = "xpt"
+optype = "op"
+"optype.numpy" = "onp"
+"optype.numpy.compat" = "npc"
 
 [tool.ruff.lint.isort]
 combine-as-imports = true

src/array_api_typing/__init__.py

@@ -1,10 +1,11 @@
 """Static typing support for the array API standard."""
 
 __all__ = (
+    "Array",
     "HasArrayNamespace",
     "__version__",
     "__version_tuple__",
 )
 
-from ._namespace import HasArrayNamespace
+from ._array import Array, HasArrayNamespace
 from ._version import version as __version__, version_tuple as __version_tuple__

src/array_api_typing/_array.py

@@ -0,0 +1,268 @@
+__all__ = (
+    "Array",
+    "BoolArray",
+    "HasArrayNamespace",
+    "NumericArray",
+)
+
+from pathlib import Path
+from types import ModuleType
+from typing import Any, Literal, Never, Protocol, Self, TypeAlias
+from typing_extensions import TypeVar
+
+import optype as op
+
+from ._utils import docstring_setter
+
+# Load docstrings from TOML file
+try:
+    import tomllib
+except ImportError:
+    import tomli as tomllib  # type: ignore[import-not-found, no-redef]
+
+_docstrings_path = Path(__file__).parent / "_array_docstrings.toml"
+with _docstrings_path.open("rb") as f:
+    _array_docstrings = tomllib.load(f)["docstrings"]
+
+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]):
+    """Protocol for classes that have an `__array_namespace__` method.
+
+    Example:
+    >>> import array_api_typing as xpt
+    >>>
+    >>> class MyArray:
+    ...     def __array_namespace__(self):
+    ...         return object()
+    >>>
+    >>> x = MyArray()
+    >>> def has_array_namespace(x: xpt.HasArrayNamespace) -> bool:
+    ...     return hasattr(x, "__array_namespace__")
+    >>> has_array_namespace(x)
+    True
+
+    """
+
+    def __array_namespace__(
+        self, /, *, api_version: Literal["2021.12"] | None = None
+    ) -> 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."""
+        ...
+
+
+class HasDevice(Protocol):
+    """Protocol for array classes that have a device attribute."""
+
+    @property
+    def device(self) -> object:  # TODO: more specific type
+        """Hardware device the array data resides on."""
+        ...
+
+
+class HasMatrixTranspose(Protocol):
+    """Protocol for array classes that have a matrix transpose attribute."""
+
+    @property
+    def mT(self) -> Self:  # noqa: N802
+        """Transpose of a matrix (or a stack of matrices).
+
+        If an array instance has fewer than two dimensions, an error should be
+        raised.
+
+        Returns:
+            Self: array whose last two dimensions (axes) are permuted in reverse
+                order relative to original array (i.e., for an array instance
+                having shape `(..., M, N)`, the returned array must have shape
+                `(..., N, M))`.  The returned array must have the same data type
+                as the original array.
+
+        """
+        ...
+
+
+class HasNDim(Protocol):
+    """Protocol for array classes that have a number of dimensions attribute."""
+
+    @property
+    def ndim(self) -> int:
+        """Number of array dimensions (axes).
+
+        Returns:
+            int: number of array dimensions (axes).
+
+        """
+        ...
+
+
+class HasShape(Protocol):
+    """Protocol for array classes that have a shape attribute."""
+
+    @property
+    def shape(self) -> tuple[int | None, ...]:
+        """Shape of the array.
+
+        Returns:
+            tuple[int | None, ...]: array dimensions. An array dimension must be None
+                if and only if a dimension is unknown.
+
+        Notes:
+            For array libraries having graph-based computational models, array
+            dimensions may be unknown due to data-dependent operations (e.g.,
+            boolean indexing; `A[:, B > 0]`) and thus cannot be statically
+            resolved without knowing array contents.
+
+        """
+        ...
+
+
+class HasSize(Protocol):
+    """Protocol for array classes that have a size attribute."""
+
+    @property
+    def size(self) -> int | None:
+        """Number of elements in an array.
+
+        Returns:
+            int | None: number of elements in an array. The returned value must
+                be `None` if and only if one or more array dimensions are
+                unknown.
+
+        Notes:
+            This must equal the product of the array's dimensions.
+
+        """
+        ...
+
+
+class HasTranspose(Protocol):
+    """Protocol for array classes that support the transpose operation."""
+
+    def T(self) -> Self:  # noqa: N802
+        """Transpose of the array.
+
+        The array instance must be two-dimensional. If the array instance is not
+        two-dimensional, an error should be raised.
+
+        Returns:
+            Self: two-dimensional array whose first and last dimensions (axes)
+                are permuted in reverse order relative to original array. The
+                returned array must have the same data type as the original
+                array.
+
+        Notes:
+            Limiting the transpose to two-dimensional arrays (matrices) deviates
+            from the NumPy et al practice of reversing all axes for arrays
+            having more than two-dimensions. This is intentional, as reversing
+            all axes was found to be problematic (e.g., conflicting with the
+            mathematical definition of a transpose which is limited to matrices;
+            not operating on batches of matrices; et cetera). In order to
+            reverse all axes, one is recommended to use the functional
+            `PermuteDims` interface found in this specification.
+
+        """
+        ...
+
+
+@docstring_setter(**_array_docstrings)
+class Array(
+    # ------ Attributes -------
+    HasDType[DType_co],
+    HasDevice,
+    HasMatrixTranspose,
+    HasNDim,
+    HasShape,
+    HasSize,
+    HasTranspose,
+    # ------ Methods -------
+    HasArrayNamespace[NS_co],
+    op.CanPosSelf,
+    op.CanNegSelf,
+    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.
+
+    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 = 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
+      allow for wrapper libraries to semi-dynamically define their own array
+      namespaces based on the wrapped array type.
+
+    """
+
+
+# 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.
+
+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.
+
+"""
+
+# 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.
+
+"""