MAINT: backports and tweaks

Author: crusaderky

docs/conf.py

@@ -54,6 +54,7 @@
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
     "array-api": ("https://data-apis.org/array-api/draft", None),
+    "dask": ("https://docs.dask.org/en/stable", None),
     "numpy": ("https://numpy.org/doc/stable", None),
     "jax": ("https://jax.readthedocs.io/en/latest", None),
 }

pixi.lock

@@ -191,10 +191,6 @@ markers = [
 
 [tool.coverage]
 run.source = ["array_api_extra"]
-report.exclude_also = [
-    '\.\.\.',
-    'if TYPE_CHECKING:',
-]
 
 # mypy
 
@@ -314,10 +310,5 @@ checks = [
     "ES01", # most docstrings do not need an extended summary
 ]
 exclude = [ # don't report on objects that match any of these regex
-    '.*test_at.*',
-    '.*test_funcs.*',
-    '.*test_testing.*',
-    '.*test_utils.*',
-    '.*test_version.*',
-    '.*test_vendor.*',
+    '.*test_*',
 ]

pyproject.toml

@@ -249,8 +249,8 @@ def _op(
             Right-hand side of the operation.
         copy : bool or None
             Whether to copy the input array. See the class docstring for details.
-        xp : array_namespace or None
-            The array namespace for the input array.
+        xp : array_namespace, optional
+            The array namespace for the input array. Default: infer.
 
         Returns
         -------

src/array_api_extra/_lib/_at.py

@@ -7,10 +7,10 @@
 from typing import cast
 
 from . import _compat
-from ._compat import is_array_api_obj, is_numpy_array
+from ._compat import array_namespace, is_array_api_obj, is_numpy_array
 from ._typing import Array
 
-__all__ = ["in1d", "mean"]
+__all__ = ["asarrays", "in1d", "is_python_scalar", "mean"]
 
 
 def in1d(
@@ -33,7 +33,7 @@ def in1d(
     https://github.com/numpy/numpy/blob/v1.26.0/numpy/lib/arraysetops.py#L524-L758
     """
     if xp is None:
-        xp = _compat.array_namespace(x1, x2)
+        xp = array_namespace(x1, x2)
 
     # This code is run to make the code significantly faster
     if x2.shape[0] < 10 * x1.shape[0] ** 0.145:
@@ -84,7 +84,7 @@ def mean(
     Complex mean, https://github.com/data-apis/array-api/issues/846.
     """
     if xp is None:
-        xp = _compat.array_namespace(x)
+        xp = array_namespace(x)
 
     if xp.isdtype(x.dtype, "complex floating"):
         x_real = xp.real(x)
@@ -124,8 +124,8 @@ def asarrays(
     ----------
     a, b : Array | int | float | complex | bool
         Input arrays or scalars. At least one must be an array.
-    xp : ModuleType
-        The standard-compatible namespace for the returned arrays.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
 
     Returns
     -------

src/array_api_extra/_lib/_utils/_helpers.py

@@ -5,6 +5,7 @@
 # To be changed to a Protocol later (see data-apis/array-api#589)
 Array = Any  # type: ignore[no-any-explicit]
 Device = Any  # type: ignore[no-any-explicit]
+DType = Any  # type: ignore[no-any-explicit]
 Index = Any  # type: ignore[no-any-explicit]
 
-__all__ = ["Array", "Device", "Index"]
+__all__ = ["Array", "DType", "Device", "Index"]

src/array_api_extra/_lib/_utils/_typing.py

@@ -18,7 +18,7 @@
 
 __all__ = ["lazy_xp_function", "patch_lazy_xp_functions"]
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     # TODO move ParamSpec outside TYPE_CHECKING
     # depends on scikit-learn abandoning Python 3.9
     # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
@@ -169,7 +169,7 @@ def xp(request, monkeypatch):
         Pytest fixture, as acquired by the test itself or by one of its fixtures.
     monkeypatch : pytest.MonkeyPatch
         Pytest fixture, as acquired by the test itself or by one of its fixtures.
-    xp : module
+    xp : array_namespace
         Array namespace to be tested.
 
     See Also

src/array_api_extra/testing.py

@@ -115,7 +115,11 @@ def xp(
     if library == Backend.NUMPY_READONLY:
         return NumPyReadOnly()  # type: ignore[return-value]  # pyright: ignore[reportReturnType]
     xp = pytest.importorskip(library.value)
+    # Possibly wrap module with array_api_compat
+    xp = array_namespace(xp.empty(0))
 
+    # On Dask and JAX, monkey-patch all functions tagged by `lazy_xp_function`
+    # in the global scope of the module containing the test function.
     patch_lazy_xp_functions(request, monkeypatch, xp=xp)
 
     if library == Backend.JAX:
@@ -124,8 +128,18 @@ def xp(
         # suppress unused-ignore to run mypy in -e lint as well as -e dev
         jax.config.update("jax_enable_x64", True)  # type: ignore[no-untyped-call,unused-ignore]
 
-    # Possibly wrap module with array_api_compat
-    return array_namespace(xp.empty(0))
+    return xp
+
+
+@pytest.fixture(params=[Backend.DASK])  # Can select the test with `pytest -k dask`
+def da(
+    request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch
+) -> ModuleType:  # numpydoc ignore=PR01,RT01
+    """Variant of the `xp` fixture that only yields dask.array."""
+    xp = pytest.importorskip("dask.array")
+    xp = array_namespace(xp.empty(0))
+    patch_lazy_xp_functions(request, monkeypatch, xp=xp)
+    return xp
 
 
 @pytest.fixture