✨ add CanArray* unop and binop protocols

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

Author: nstarman

pyproject.toml

@@ -29,6 +29,7 @@ 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]
@@ -123,9 +124,12 @@ 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`
 ]
 
 [tool.ruff.lint.pylint]

src/array_api_typing/_array.py

@@ -1,13 +1,32 @@
 __all__ = (
     "Array",
+    "BoolArray",
     "HasArrayNamespace",
+    "NumericArray",
 )
 
+from pathlib import Path
 from types import ModuleType
-from typing import Literal, Protocol
+from typing import Literal, Never, Protocol, 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)
+T_contra = TypeVar("T_contra", contravariant=True)
+R_co = TypeVar("R_co", covariant=True, default=Never)
 
 
 class HasArrayNamespace(Protocol[NS_co]):
@@ -33,8 +52,37 @@ def __array_namespace__(
     ) -> NS_co: ...
 
 
+@docstring_setter(**_array_docstrings)
 class Array(
     HasArrayNamespace[NS_co],
-    Protocol[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],
 ):
     """Array API specification for array object attributes and methods."""
+
+
+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
+`bool`, allowing for `bool` objects to be added, subtracted, multiplied, etc. to
+the array object.
+
+"""
+
+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
+| int`, allowing for `float | int` objects to be added, subtracted, multiplied,
+etc. to the array object.
+
+"""

src/array_api_typing/_array_docstrings.toml

@@ -0,0 +1,151 @@
+[docstrings]
+__pos__ = '''
+Evaluates `+self_i` for each element of an array instance.
+
+Returns:
+    Self: An array containing the evaluated result for each element.
+        The returned array must have the same data type as `self`.
+
+See Also:
+    array_api_typing.Positive
+
+'''
+
+__neg__ = '''
+Evaluates `-self_i` for each element of an array instance.
+
+Returns:
+    Self: an array containing the evaluated result for each element in
+        `self`. The returned array must have a data type determined by Type
+        Promotion Rules.
+
+See Also:
+    array_api_typing.Negative
+
+'''
+
+__add__ = '''
+Calculates the sum for each element of an array instance with the respective
+element of the array `other`.
+
+Args:
+    other: addend array. Must be compatible with `self` (see
+        Broadcasting). Should have a numeric data type.
+
+Returns:
+    Self: an array containing the element-wise sums. The returned array
+        must have a data type determined by Type Promotion Rules.
+
+See Also:
+    array_api_typing.Add
+
+'''
+
+__sub__ = '''
+Calculates the difference for each element of an array instance with the
+respective element of the array other.
+
+The result of `self_i - other_i` must be the same as `self_i +
+(-other_i)` and must be governed by the same floating-point rules as
+addition (see `CanArrayAdd`).
+
+Args:
+    other: subtrahend array. Must be compatible with `self` (see
+        Broadcasting). Should have a numeric data type.
+
+Returns:
+    Self: an array containing the element-wise differences. The returned
+        array must have a data type determined by Type Promotion Rules.
+
+See Also:
+    array_api_typing.Subtract
+
+'''
+
+__mul__ = '''
+Calculates the product for each element of an array instance with the
+respective element of the array `other`.
+
+Args:
+    other: multiplicand array. Must be compatible with `self` (see
+        Broadcasting). Should have a numeric data type.
+
+Returns:
+    Self: an array containing the element-wise products. The returned
+        array must have a data type determined by Type Promotion Rules.
+
+See Also:
+    array_api_typing.Multiply
+
+'''
+
+__truediv__ = '''
+Evaluates `self_i / other_i` for each element of an array instance with the
+respective element of the array `other`.
+
+Args:
+    other: Must be compatible with `self` (see Broadcasting). Should have a
+        numeric data type.
+
+Returns:
+    Self: an array containing the element-wise results. The returned array
+        should have a floating-point data type determined by Type Promotion
+        Rules.
+
+See Also:
+    array_api_typing.TrueDiv
+
+'''
+
+__floordiv__ = '''
+Evaluates `self_i // other_i` for each element of an array instance with the
+respective element of the array `other`.
+
+Args:
+    other: Must be compatible with `self` (see Broadcasting). Should have a
+        numeric data type.
+
+Returns:
+    Self: an array containing the element-wise results. The returned array
+        must have a data type determined by Type Promotion Rules.
+
+See Also:
+    array_api_typing.FloorDiv
+
+'''
+
+__mod__ = '''
+Evaluates `self_i % other_i` for each element of an array instance with the
+respective element of the array `other`.
+
+Args:
+    other: Must be compatible with `self` (see Broadcasting). Should have a
+        numeric data type.
+
+Returns:
+    Self: an array containing the element-wise results. Each element-wise
+        result must have the same sign as the respective element `other_i`.
+        The returned array must have a floating-point data type determined
+        by Type Promotion Rules.
+
+See Also:
+    array_api_typing.Remainder
+
+'''
+
+__pow__ = '''
+Calculates an implementation-dependent approximation of exponentiation by
+raising each element (the base) of an array instance to the power of
+`other_i` (the exponent), where `other_i` is the corresponding element of
+the array `other`.
+
+Args:
+    other: array whose elements correspond to the exponentiation exponent.
+        Must be compatible with `self` (see Broadcasting). Should have a
+        numeric data type.
+
+Returns:
+    Self: an array containing the element-wise results. The returned array
+        must have a data type determined by Type Promotion Rules.
+
+'''

src/array_api_typing/_namespace.py

@@ -0,0 +1,64 @@
+"""Utility functions."""
+
+from collections.abc import Callable
+from enum import Enum, auto
+from typing import Literal, TypeVar
+
+ClassT = TypeVar("ClassT")
+DocstringTypes = str | None
+
+
+class _Sentinel(Enum):
+    SKIP = auto()
+
+
+def set_docstrings(
+    obj: type[ClassT],
+    main: DocstringTypes | Literal[_Sentinel.SKIP] = _Sentinel.SKIP,
+    /,
+    **method_docs: DocstringTypes,
+) -> type[ClassT]:
+    """Set the docstring for a class and its methods.
+
+    Args:
+        obj: The class to set the docstring for.
+        main: The main docstring for the class. If not provided, the
+            class docstring will not be modified.
+        method_docs: A mapping of method names to their docstrings. If a method
+            is not provided, its docstring will not be modified.
+
+    Returns:
+        The class with updated docstrings.
+
+    """
+    if main is not _Sentinel.SKIP:
+        obj.__doc__ = main
+
+    for name, doc in method_docs.items():
+        method = getattr(obj, name)
+        method.__doc__ = doc
+    return obj
+
+
+def docstring_setter(
+    main: DocstringTypes | Literal[_Sentinel.SKIP] = _Sentinel.SKIP,
+    /,
+    **method_docs: DocstringTypes,
+) -> Callable[[type[ClassT]], type[ClassT]]:
+    """Decorator to set docstrings for a class and its methods.
+
+    Args:
+        main: The main docstring for the class. If not provided, the
+            class docstring will not be modified.
+        method_docs: A mapping of method names to their docstrings. If a method
+            is not provided, its docstring will not be modified.
+
+    Returns:
+        A decorator that sets the docstrings for the class and its methods.
+
+    """
+
+    def decorator(cls: type[ClassT]) -> type[ClassT]:
+        return set_docstrings(cls, main, **method_docs)
+
+    return decorator

src/array_api_typing/_utils.py

@@ -1,12 +1,53 @@
-from typing import Any
+# Test array_api_typing with numpy < 2.0
 
-# requires numpy < 2
-import numpy.array_api as np  # type: ignore[import-not-found]
+from typing import Any, Never, TypeAlias
+
+import numpy.array_api as np  # type: ignore[import-not-found, unused-ignore]
+from numpy import dtype, floating, integer
 
 import array_api_typing as xpt
+from array_api_typing._array import BoolArray, NumericArray
+
+F: TypeAlias = floating[Any]
+F32: TypeAlias = dtype  # Note: np.array_api uses dtype objects.
+I: TypeAlias = integer[Any]
+I32: TypeAlias = dtype  # Note: np.array_api uses dtype objects.
+
+# Define an NDArray against which we can test the protocols
+nparr = np.eye(2)
+nparr_i32 = np.asarray([1], dtype=np.int32)
+nparr_f32 = np.asarray([1.0], dtype=np.float32)
+nparr_b = np.asarray([True], dtype=bool)
+
+# =========================================================
+# Ensure that `np.ndarray` instances are assignable to `xpt.HasArrayNamespace`
+
+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.Array`.
+
+# Generic Array type
+arr_array: xpt.Array[Never] = 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
+
+# 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
+
+# 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
 
-###
-# Ensure that `np.ndarray` instances are assignable to `xpt.HasArrayNamespace`.
+# =========================================================
+# Check np.ndarray against BoolArray and NumericArray type aliases
 
-arr = np.eye(2)
-arr_namespace: xpt.HasArrayNamespace[Any] = arr
+boolarray: BoolArray = nparr_b
+numericarray: NumericArray = nparr_f32