ENH Test tools for jax.jit and dask

Author: crusaderky

docs/api-reference.md

@@ -17,3 +17,15 @@
     setdiff1d
     sinc
 ```
+
+## Testing utilities
+
+```{eval-rst}
+.. currentmodule:: array_api_extra.testing
+.. autosummary::
+    :nosignatures:
+    :toctree: generated
+
+    lazy_xp_function
+    patch_lazy_xp_functions
+```

pixi.lock

@@ -98,6 +98,8 @@ furo = ">=2023.08.17"
 myst-parser = ">=0.13"
 sphinx-copybutton = "*"
 sphinx-autodoc-typehints = "*"
+# Needed to import parsed modules with autodoc
+pytest = "*"
 
 [tool.pixi.feature.docs.tasks]
 docs = { cmd = "sphinx-build . build/", cwd = "docs" }
@@ -180,8 +182,10 @@ markers = ["skip_xp_backend(library, *, reason=None): Skip test for a specific b
 
 [tool.coverage]
 run.source = ["array_api_extra"]
-report.exclude_also = ['\.\.\.']
-
+report.exclude_also = [
+    '\.\.\.',
+    'if TYPE_CHECKING:',
+]
 
 # mypy
 
@@ -221,6 +225,8 @@ reportMissingImports = false
 reportMissingTypeStubs = false
 # false positives for input validation
 reportUnreachable = false
+# ruff handles this
+reportUnusedParameter = false
 
 executionEnvironments = [
     { root = "tests", reportPrivateUsage = false },
@@ -282,7 +288,10 @@ messages_control.disable = [
     "design",                     # ignore heavily opinionated design checks
     "fixme",                      # allow FIXME comments
     "line-too-long",              # ruff handles this
+    "unused-argument",            # ruff handles this
     "missing-function-docstring", # numpydoc handles this
+    "import-error",               # mypy handles this
+    "import-outside-toplevel",    # optional dependencies
 ]
 
 

pyproject.toml

@@ -6,6 +6,7 @@
 import operator
 from collections.abc import Callable
 from enum import Enum
+from numbers import Number
 from types import ModuleType
 from typing import ClassVar, cast
 
@@ -188,7 +189,7 @@ def __getitem__(self, idx: Index, /) -> at:  # numpydoc ignore=PR01,RT01
     def _update_common(
         self,
         at_op: _AtOp,
-        y: Array,
+        y: Array | Number,
         /,
         copy: bool | None,
         xp: ModuleType | None,
@@ -253,7 +254,7 @@ def _update_common(
 
     def set(
         self,
-        y: Array,
+        y: Array | Number,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -269,8 +270,8 @@ def set(
     def _iop(
         self,
         at_op: _AtOp,
-        elwise_op: Callable[[Array, Array], Array],
-        y: Array,
+        elwise_op: Callable[[Array, Array | Number], Array],
+        y: Array | Number,
         /,
         copy: bool | None,
         xp: ModuleType | None,
@@ -294,7 +295,7 @@ def _iop(
 
     def add(
         self,
-        y: Array,
+        y: Array | Number,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -308,7 +309,7 @@ def add(
 
     def subtract(
         self,
-        y: Array,
+        y: Array | Number,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -318,7 +319,7 @@ def subtract(
 
     def multiply(
         self,
-        y: Array,
+        y: Array | Number,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -328,7 +329,7 @@ def multiply(
 
     def divide(
         self,
-        y: Array,
+        y: Array | Number,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -338,7 +339,7 @@ def divide(
 
     def power(
         self,
-        y: Array,
+        y: Array | Number,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -348,7 +349,7 @@ def power(
 
     def min(
         self,
-        y: Array,
+        y: Array | Number,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -361,7 +362,7 @@ def min(
 
     def max(
         self,
-        y: Array,
+        y: Array | Number,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,

src/array_api_extra/_lib/_at.py

@@ -2,6 +2,7 @@
 Testing utilities.
 
 Note that this is private API; don't expect it to be stable.
+See also ..testing for public testing utilities.
 """
 
 from types import ModuleType