✨: add CanArrayX protocols

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,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]
@@ -137,10 +143,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,88 @@
+__all__ = (
+    "Array",
+    "BoolArray",
+    "HasArrayNamespace",
+    "NumericArray",
+)
+
+from pathlib import Path
+from types import ModuleType
+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]):
+    """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: ...
+
+
+@docstring_setter(**_array_docstrings)
+class Array(
+    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],
+):
+    """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

@@ -1,30 +0,0 @@
-__all__ = ("HasArrayNamespace",)
-
-from types import ModuleType
-from typing import Literal, Protocol
-from typing_extensions import TypeVar
-
-T_co = TypeVar("T_co", covariant=True, default=ModuleType)
-
-
-class HasArrayNamespace(Protocol[T_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
-    ) -> T_co: ...

src/array_api_typing/_utils.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