TST: Run all tests on read-only numpy arrays

Author: crusaderky

tests/conftest.py

@@ -1,14 +1,24 @@
 """Pytest fixtures."""
 
+from __future__ import annotations
+
+from collections.abc import Callable
 from enum import Enum
-from typing import cast
+from functools import wraps
+from typing import ParamSpec, TypeVar, cast
 
+import numpy as np
 import pytest
 
 from array_api_extra._lib._compat import array_namespace
 from array_api_extra._lib._compat import device as get_device
 from array_api_extra._lib._typing import Device, ModuleType
 
+T = TypeVar("T")
+P = ParamSpec("P")
+
+np_compat = array_namespace(np.empty(0))
+
 
 class Library(Enum):
     """All array libraries explicitly tested by array-api-extra."""
@@ -50,6 +60,56 @@ def library(request: pytest.FixtureRequest) -> Library:  # numpydoc ignore=PR01,
     return elem
 
 
+class NumPyReadOnly:
+    """
+    Variant of array_api_compat.numpy producing read-only arrays.
+
+    Note that this is not a full read-only Array API library. Notably,
+    array_namespace(x) returns array_api_compat.numpy, and as a consequence array
+    creation functions invoked internally by the tested functions will return
+    writeable arrays, as long as you don't explicitly pass xp=xp.
+    For this reason, tests that do pass xp=xp may misbehave and should be skipped
+    for NUMPY_READONLY.
+    """
+
+    def __getattr__(self, name: str) -> object:  # numpydoc ignore=PR01,RT01
+        """Wrap all functions that return arrays to make their output read-only."""
+        func = getattr(np_compat, name)
+        if not callable(func) or isinstance(func, type):
+            return func
+        return self._wrap(func)
+
+    @staticmethod
+    def _wrap(func: Callable[P, T]) -> Callable[P, T]:  # numpydoc ignore=PR01,RT01
+        """Wrap func to make all np.ndarrays it returns read-only."""
+
+        def as_readonly(o: T, seen: set[int]) -> T:  # numpydoc ignore=PR01,RT01
+            """Unset the writeable flag in o."""
+            if id(o) in seen:
+                return o
+            seen.add(id(o))
+
+            try:
+                # Don't use is_numpy_array(o), as it includes np.generic
+                if isinstance(o, np.ndarray):
+                    o.flags.writeable = False
+            except TypeError:
+                # Cannot interpret as a data type
+                return o
+
+            # This works with namedtuples too
+            if isinstance(o, tuple | list):
+                return type(o)(*(as_readonly(i, seen) for i in o))  # type: ignore[arg-type,return-value] # pyright: ignore[reportArgumentType,reportUnknownArgumentType]
+
+            return o
+
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
+            return as_readonly(func(*args, **kwargs), seen=set())
+
+        return wrapper
+
+
 @pytest.fixture
 def xp(library: Library) -> ModuleType:  # numpydoc ignore=PR01,RT03
     """
@@ -59,8 +119,9 @@ def xp(library: Library) -> ModuleType:  # numpydoc ignore=PR01,RT03
     -------
     The current array namespace.
     """
-    name = "numpy" if library == Library.NUMPY_READONLY else library.value
-    xp = pytest.importorskip(name)
+    if library == Library.NUMPY_READONLY:
+        return NumPyReadOnly()  # type: ignore[return-value]  # pyright: ignore[reportReturnType]
+    xp = pytest.importorskip(library.value)
     if library == Library.JAX_NUMPY:
         import jax  # type: ignore[import-not-found]  # pyright: ignore[reportMissingImports]
 

tests/test_at.py

@@ -6,7 +6,6 @@
 import pytest
 from array_api_compat import (  # type: ignore[import-untyped]  # pyright: ignore[reportMissingTypeStubs]
     array_namespace,
-    is_pydata_sparse_array,
     is_writeable_array,
 )
 
@@ -18,14 +17,6 @@
 from .conftest import Library
 
 
-@pytest.fixture
-def array(library: Library, xp: ModuleType) -> Array:
-    x = xp.asarray([10.0, 20.0, 30.0])
-    if library == Library.NUMPY_READONLY:
-        x.flags.writeable = False
-    return x
-
-
 @contextmanager
 def assert_copy(array: Array, copy: bool | None) -> Generator[None, None, None]:
     if copy is False and not is_writeable_array(array):
@@ -42,6 +33,9 @@ def assert_copy(array: Array, copy: bool | None) -> Generator[None, None, None]:
     xp_assert_equal(xp.all(array == array_orig), xp.asarray(copy))
 
 
+@pytest.mark.skip_xp_backend(
+    Library.SPARSE, reason="read-only library without .at support"
+)
 @pytest.mark.parametrize(
     ("kwargs", "expect_copy"),
     [
@@ -66,15 +60,13 @@ def assert_copy(array: Array, copy: bool | None) -> Generator[None, None, None]:
 )
 def test_update_ops(
     xp: ModuleType,
-    array: Array,
     kwargs: dict[str, bool | None],
     expect_copy: bool | None,
     op: _AtOp,
     arg: float,
     expect: list[float],
 ):
-    if is_pydata_sparse_array(array):
-        pytest.skip("at() does not support updates on sparse arrays")
+    array = xp.asarray([10.0, 20.0, 30.0])
 
     with assert_copy(array, expect_copy):
         func = cast(Callable[..., Array], getattr(at(array)[1:], op.value))  # type: ignore[no-any-explicit]

tests/test_funcs.py

@@ -136,6 +136,7 @@ def test_device(self, xp: ModuleType, device: Device):
         x = xp.asarray([1, 2, 3], device=device)
         assert get_device(cov(x)) == device
 
+    @pytest.mark.skip_xp_backend(Library.NUMPY_READONLY)
     def test_xp(self, xp: ModuleType):
         xp_assert_close(
             cov(xp.asarray([[0.0, 2.0], [1.0, 1.0], [2.0, 0.0]]).T, xp=xp),
@@ -366,6 +367,7 @@ def test_device(self, xp: ModuleType, device: Device):
         x2 = xp.asarray([2, 3, 4], device=device)
         assert get_device(setdiff1d(x1, x2)) == device
 
+    @pytest.mark.skip_xp_backend(Library.NUMPY_READONLY)
     def test_xp(self, xp: ModuleType):
         x1 = xp.asarray([3, 8, 20])
         x2 = xp.asarray([2, 3, 4])

tests/test_utils.py

@@ -1,9 +1,8 @@
-import numpy as np
 import pytest
 
 from array_api_extra._lib._compat import device as get_device
 from array_api_extra._lib._testing import xp_assert_equal
-from array_api_extra._lib._typing import Array, Device, ModuleType
+from array_api_extra._lib._typing import Device, ModuleType
 from array_api_extra._lib._utils import in1d
 
 from .conftest import Library
@@ -15,10 +14,10 @@ class TestIn1D:
     @pytest.mark.skip_xp_backend(Library.DASK_ARRAY, reason="no argsort")
     @pytest.mark.skip_xp_backend(Library.SPARSE, reason="no unique_inverse, no device")
     # cover both code paths
-    @pytest.mark.parametrize("x2", [np.arange(9), np.arange(15)])
-    def test_no_invert_assume_unique(self, xp: ModuleType, x2: Array):
+    @pytest.mark.parametrize("n", [9, 15])
+    def test_no_invert_assume_unique(self, xp: ModuleType, n: int):
         x1 = xp.asarray([3, 8, 20])
-        x2 = xp.asarray(x2)
+        x2 = xp.arange(n)
         expected = xp.asarray([True, True, False])
         actual = in1d(x1, x2)
         xp_assert_equal(actual, expected)
@@ -29,6 +28,7 @@ def test_device(self, xp: ModuleType, device: Device):
         x2 = xp.asarray([2, 3, 4], device=device)
         assert get_device(in1d(x1, x2)) == device
 
+    @pytest.mark.skip_xp_backend(Library.NUMPY_READONLY)
     @pytest.mark.skip_xp_backend(Library.SPARSE, reason="no arange, no device")
     def test_xp(self, xp: ModuleType):
         x1 = xp.asarray([1, 6])