Add an Array Protocol & improve static typing support

.pre-commit-config.yaml

@@ -48,3 +48,29 @@ repos:
     rev: 23.7.0
     hooks:
       - id: black
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: "v1.0.0"
+    hooks:
+      - id: mypy
+        additional_dependencies: [typing_extensions>=4.4.0]
+        args:
+          - --ignore-missing-imports
+          - --config=pyproject.toml
+        files: ".*(_draft.*)$"
+        exclude: |
+          (?x)^(
+              .*creation_functions.py|
+              .*data_type_functions.py|
+              .*elementwise_functions.py|
+              .*fft.py|
+              .*indexing_functions.py|
+              .*linalg.py|
+              .*linear_algebra_functions.py|
+              .*manipulation_functions.py|
+              .*searching_functions.py|
+              .*set_functions.py|
+              .*sorting_functions.py|
+              .*statistical_functions.py|
+              .*utility_functions.py|
+          )$

pyproject.toml

@@ -30,5 +30,15 @@ doc = [
 requires = ["setuptools"]
 build-backend = "setuptools.build_meta"
 
+
 [tool.black]
 line-length = 88
+
+
+[tool.mypy]
+python_version = "3.9"
+mypy_path = "$MYPY_CONFIG_FILE_DIR/src/array_api_stubs/_draft/"
+files = [
+    "src/array_api_stubs/_draft/**/*.py"
+]
+follow_imports = "silent"

spec/draft/API_specification/array_object.rst

@@ -30,47 +30,47 @@ Arithmetic Operators
 
 A conforming implementation of the array API standard must provide and support an array object supporting the following Python arithmetic operators.
 
--   ``+x``: :meth:`.array.__pos__`
+-   ``+x``: :meth:`.Array.__pos__`
 
     -   `operator.pos(x) <https://docs.python.org/3/library/operator.html#operator.pos>`_
     -   `operator.__pos__(x) <https://docs.python.org/3/library/operator.html#operator.__pos__>`_
 
--   `-x`: :meth:`.array.__neg__`
+-   `-x`: :meth:`.Array.__neg__`
 
     -   `operator.neg(x) <https://docs.python.org/3/library/operator.html#operator.neg>`_
     -   `operator.__neg__(x) <https://docs.python.org/3/library/operator.html#operator.__neg__>`_
 
--   `x1 + x2`: :meth:`.array.__add__`
+-   `x1 + x2`: :meth:`.Array.__add__`
 
     -   `operator.add(x1, x2) <https://docs.python.org/3/library/operator.html#operator.add>`_
     -   `operator.__add__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__add__>`_
 
--   `x1 - x2`: :meth:`.array.__sub__`
+-   `x1 - x2`: :meth:`.Array.__sub__`
 
     -   `operator.sub(x1, x2) <https://docs.python.org/3/library/operator.html#operator.sub>`_
     -   `operator.__sub__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__sub__>`_
 
--   `x1 * x2`: :meth:`.array.__mul__`
+-   `x1 * x2`: :meth:`.Array.__mul__`
 
     -   `operator.mul(x1, x2) <https://docs.python.org/3/library/operator.html#operator.mul>`_
     -   `operator.__mul__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__mul__>`_
 
--   `x1 / x2`: :meth:`.array.__truediv__`
+-   `x1 / x2`: :meth:`.Array.__truediv__`
 
     -   `operator.truediv(x1,x2) <https://docs.python.org/3/library/operator.html#operator.truediv>`_
     -   `operator.__truediv__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__truediv__>`_
 
--   `x1 // x2`: :meth:`.array.__floordiv__`
+-   `x1 // x2`: :meth:`.Array.__floordiv__`
 
     -   `operator.floordiv(x1, x2) <https://docs.python.org/3/library/operator.html#operator.floordiv>`_
     -   `operator.__floordiv__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__floordiv__>`_
 
--   `x1 % x2`: :meth:`.array.__mod__`
+-   `x1 % x2`: :meth:`.Array.__mod__`
 
     -   `operator.mod(x1, x2) <https://docs.python.org/3/library/operator.html#operator.mod>`_
     -   `operator.__mod__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__mod__>`_
 
--   `x1 ** x2`: :meth:`.array.__pow__`
+-   `x1 ** x2`: :meth:`.Array.__pow__`
 
     -   `operator.pow(x1, x2) <https://docs.python.org/3/library/operator.html#operator.pow>`_
     -   `operator.__pow__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__pow__>`_
@@ -82,7 +82,7 @@ Array Operators
 
 A conforming implementation of the array API standard must provide and support an array object supporting the following Python array operators.
 
--   `x1 @ x2`: :meth:`.array.__matmul__`
+-   `x1 @ x2`: :meth:`.Array.__matmul__`
 
     -   `operator.matmul(x1, x2) <https://docs.python.org/3/library/operator.html#operator.matmul>`_
     -   `operator.__matmul__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__matmul__>`_
@@ -94,34 +94,34 @@ Bitwise Operators
 
 A conforming implementation of the array API standard must provide and support an array object supporting the following Python bitwise operators.
 
--   `~x`: :meth:`.array.__invert__`
+-   `~x`: :meth:`.Array.__invert__`
 
     -   `operator.inv(x) <https://docs.python.org/3/library/operator.html#operator.inv>`_
     -   `operator.invert(x) <https://docs.python.org/3/library/operator.html#operator.invert>`_
     -   `operator.__inv__(x) <https://docs.python.org/3/library/operator.html#operator.__inv__>`_
     -   `operator.__invert__(x) <https://docs.python.org/3/library/operator.html#operator.__invert__>`_
 
--   `x1 & x2`: :meth:`.array.__and__`
+-   `x1 & x2`: :meth:`.Array.__and__`
 
     -   `operator.and(x1, x2) <https://docs.python.org/3/library/operator.html#operator.and>`_
     -   `operator.__and__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__and__>`_
 
--   `x1 | x2`: :meth:`.array.__or__`
+-   `x1 | x2`: :meth:`.Array.__or__`
 
     -   `operator.or(x1, x2) <https://docs.python.org/3/library/operator.html#operator.or>`_
     -   `operator.__or__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__or__>`_
 
--   `x1 ^ x2`: :meth:`.array.__xor__`
+-   `x1 ^ x2`: :meth:`.Array.__xor__`
 
     -   `operator.xor(x1, x2) <https://docs.python.org/3/library/operator.html#operator.xor>`_
     -   `operator.__xor__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__xor__>`_
 
--   `x1 << x2`: :meth:`.array.__lshift__`
+-   `x1 << x2`: :meth:`.Array.__lshift__`
 
     -   `operator.lshift(x1, x2) <https://docs.python.org/3/library/operator.html#operator.lshift>`_
     -   `operator.__lshift__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__lshift__>`_
 
--   `x1 >> x2`: :meth:`.array.__rshift__`
+-   `x1 >> x2`: :meth:`.Array.__rshift__`
 
     -   `operator.rshift(x1, x2) <https://docs.python.org/3/library/operator.html#operator.rshift>`_
     -   `operator.__rshift__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__rshift__>`_
@@ -133,37 +133,37 @@ Comparison Operators
 
 A conforming implementation of the array API standard must provide and support an array object supporting the following Python comparison operators.
 
--   `x1 < x2`: :meth:`.array.__lt__`
+-   `x1 < x2`: :meth:`.Array.__lt__`
 
     -   `operator.lt(x1, x2) <https://docs.python.org/3/library/operator.html#operator.lt>`_
     -   `operator.__lt__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__lt__>`_
 
--   `x1 <= x2`: :meth:`.array.__le__`
+-   `x1 <= x2`: :meth:`.Array.__le__`
 
     -   `operator.le(x1, x2) <https://docs.python.org/3/library/operator.html#operator.le>`_
     -   `operator.__le__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__le__>`_
 
--   `x1 > x2`: :meth:`.array.__gt__`
+-   `x1 > x2`: :meth:`.Array.__gt__`
 
     -   `operator.gt(x1, x2) <https://docs.python.org/3/library/operator.html#operator.gt>`_
     -   `operator.__gt__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__gt__>`_
 
--   `x1 >= x2`: :meth:`.array.__ge__`
+-   `x1 >= x2`: :meth:`.Array.__ge__`
 
     -   `operator.ge(x1, x2) <https://docs.python.org/3/library/operator.html#operator.ge>`_
     -   `operator.__ge__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__ge__>`_
 
--   `x1 == x2`: :meth:`.array.__eq__`
+-   `x1 == x2`: :meth:`.Array.__eq__`
 
     -   `operator.eq(x1, x2) <https://docs.python.org/3/library/operator.html#operator.eq>`_
     -   `operator.__eq__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__eq__>`_
 
--   `x1 != x2`: :meth:`.array.__ne__`
+-   `x1 != x2`: :meth:`.Array.__ne__`
 
     -   `operator.ne(x1, x2) <https://docs.python.org/3/library/operator.html#operator.ne>`_
     -   `operator.__ne__(x1, x2) <https://docs.python.org/3/library/operator.html#operator.__ne__>`_
 
-:meth:`.array.__lt__`, :meth:`.array.__le__`, :meth:`.array.__gt__`, :meth:`.array.__ge__` are only defined for arrays having real-valued data types. Other comparison operators should be defined for arrays having any data type.
+:meth:`.Array.__lt__`, :meth:`.Array.__le__`, :meth:`.Array.__gt__`, :meth:`.Array.__ge__` are only defined for arrays having real-valued data types. Other comparison operators should be defined for arrays having any data type.
 For backward compatibility, conforming implementations may support complex numbers; however, inequality comparison of complex numbers is unspecified and thus implementation-dependent (see :ref:`complex-number-ordering`).
 
 In-place Operators
@@ -252,13 +252,13 @@ Attributes
    :toctree: generated
    :template: property.rst
 
-   array.dtype
-   array.device
-   array.mT
-   array.ndim
-   array.shape
-   array.size
-   array.T
+   Array.dtype
+   Array.device
+   Array.mT
+   Array.ndim
+   Array.shape
+   Array.size
+   Array.T
 
 -------------------------------------------------
 
@@ -272,37 +272,37 @@ Methods
    :toctree: generated
    :template: property.rst
 
-   array.__abs__
-   array.__add__
-   array.__and__
-   array.__array_namespace__
-   array.__bool__
-   array.__complex__
-   array.__dlpack__
-   array.__dlpack_device__
-   array.__eq__
-   array.__float__
-   array.__floordiv__
-   array.__ge__
-   array.__getitem__
-   array.__gt__
-   array.__index__
-   array.__int__
-   array.__invert__
-   array.__le__
-   array.__lshift__
-   array.__lt__
-   array.__matmul__
-   array.__mod__
-   array.__mul__
-   array.__ne__
-   array.__neg__
-   array.__or__
-   array.__pos__
-   array.__pow__
-   array.__rshift__
-   array.__setitem__
-   array.__sub__
-   array.__truediv__
-   array.__xor__
-   array.to_device
+   Array.__abs__
+   Array.__add__
+   Array.__and__
+   Array.__array_namespace__
+   Array.__bool__
+   Array.__complex__
+   Array.__dlpack__
+   Array.__dlpack_device__
+   Array.__eq__
+   Array.__float__
+   Array.__floordiv__
+   Array.__ge__
+   Array.__getitem__
+   Array.__gt__
+   Array.__index__
+   Array.__int__
+   Array.__invert__
+   Array.__le__
+   Array.__lshift__
+   Array.__lt__
+   Array.__matmul__
+   Array.__mod__
+   Array.__mul__
+   Array.__ne__
+   Array.__neg__
+   Array.__or__
+   Array.__pos__
+   Array.__pow__
+   Array.__rshift__
+   Array.__setitem__
+   Array.__sub__
+   Array.__truediv__
+   Array.__xor__
+   Array.to_device

spec/draft/purpose_and_scope.md

@@ -317,7 +317,7 @@ namespace (e.g. `import package_name.array_api`). This has two issues though:
 
 To address both issues, a uniform way must be provided by a conforming
 implementation to access the API namespace, namely a [method on the array
-object](array.__array_namespace__):
+object](Array.__array_namespace__):
 
 ```
 xp = x.__array_namespace__()

src/_array_api_conf.py

@@ -62,12 +62,16 @@
     ("py:obj", "typing.Union[int, float, typing.Literal[inf, - inf]]"),
     ("py:class", "int | float | ~typing.Literal[inf, -inf]"),
     ("py:class", "enum.Enum"),
+    ("py:class", "Enum"),
     ("py:class", "ellipsis"),
 ]
 nitpick_ignore_regex = [
     ("py:class", ".*array"),
+    ("py:class", ".*Array"),
     ("py:class", ".*device"),
+    ("py:class", ".*Device"),
     ("py:class", ".*dtype"),
+    ("py:class", ".*Self"),
     ("py:class", ".*NestedSequence"),
     ("py:class", ".*SupportsBufferProtocol"),
     ("py:class", ".*PyCapsule"),
@@ -84,6 +88,7 @@
     "array": "array",
     "Device": "device",
     "Dtype": "dtype",
+    "DType": "dtype",
 }
 
 # Make autosummary show the signatures of functions in the tables using actual

src/array_api_stubs/_draft/_types.py

@@ -33,6 +33,7 @@
 
 from dataclasses import dataclass
 from typing import (
+    TYPE_CHECKING,
     Any,
     List,
     Literal,
@@ -46,9 +47,22 @@
 )
 from enum import Enum
 
-array = TypeVar("array")
-device = TypeVar("device")
-dtype = TypeVar("dtype")
+
+if TYPE_CHECKING:
+    from .array_object import Array
+    from .data_types import DType
+
+
+class Device(Protocol):
+    """Protocol for device objects."""
+
+    def __eq__(self, value: Any) -> bool:
+        ...
+
+
+array = TypeVar("array", bound="Array")
+device = TypeVar("device", bound=Device)
+dtype = TypeVar("dtype", bound="DType")
 SupportsDLPack = TypeVar("SupportsDLPack")
 SupportsBufferProtocol = TypeVar("SupportsBufferProtocol")
 PyCapsule = TypeVar("PyCapsule")
@@ -66,7 +80,7 @@ class finfo_object:
     max: float
     min: float
     smallest_normal: float
-    dtype: dtype
+    dtype: DType
 
 
 @dataclass
@@ -76,7 +90,7 @@ class iinfo_object:
     bits: int
     max: int
     min: int
-    dtype: dtype
+    dtype: DType
 
 
 _T_co = TypeVar("_T_co", covariant=True)
@@ -96,17 +110,17 @@ class Info(Protocol):
     def capabilities(self) -> Capabilities:
         ...
 
-    def default_device(self) -> device:
+    def default_device(self) -> Device:
         ...
 
-    def default_dtypes(self, *, device: Optional[device]) -> DefaultDataTypes:
+    def default_dtypes(self, *, device: Optional[Device]) -> DefaultDataTypes:
         ...
 
-    def devices(self) -> List[device]:
+    def devices(self) -> List[Device]:
         ...
 
     def dtypes(
-        self, *, device: Optional[device], kind: Optional[Union[str, Tuple[str, ...]]]
+        self, *, device: Optional[Device], kind: Optional[Union[str, Tuple[str, ...]]]
     ) -> DataTypes:
         ...