at().get()

Author: crusaderky

src/array_api_extra/_funcs.py

@@ -11,7 +11,10 @@
 from ._lib import _utils
 from ._lib._compat import (
     array_namespace,
+    is_array_api_obj,
+    is_dask_array,
     is_jax_array,
+    is_pydata_sparse_array,
     is_writeable_array,
 )
 
@@ -686,6 +689,73 @@ def __getitem__(self, idx: Index, /) -> at:
         self._idx = idx
         return self
 
+    def _check_args(self, /, copy: bool | None) -> None:
+        if self._idx is _undef:
+            msg = (
+                "Index has not been set.\n"
+                "Usage: either\n"
+                "    at(x, idx).set(value)\n"
+                "or\n"
+                "    at(x)[idx].set(value)\n"
+                "(same for all other methods)."
+            )
+            raise TypeError(msg)
+
+        if copy not in (True, False, None):
+            msg = f"copy must be True, False, or None; got {copy!r}"  # pyright: ignore[reportUnreachable]
+            raise ValueError(msg)
+
+    def get(
+        self,
+        /,
+        copy: bool | None = True,
+        xp: ModuleType | None = None,
+    ) -> Array:
+        """Return ``xp.asarray(x[idx])``. In addition to plain ``__getitem__``,
+        this allows ensuring that the output is either a copy or a view
+        """
+        self._check_args(copy=copy)
+        x = self._x
+
+        if copy is False:
+            if is_array_api_obj(self._idx):
+                # Boolean index. Note that the array API spec
+                # https://data-apis.org/array-api/latest/API_specification/indexing.html
+                # does not allow for list, tuple, and tuples of slices plus one or more
+                # one-dimensional array indices, although many backends support them.
+                # So this check will encounter a lot of false negatives in real life,
+                # which can be caught by testing the user code vs. array-api-strict.
+                msg = "get() with an array index always returns a copy"
+                raise ValueError(msg)
+
+            # Prevent scalar indices together with copy=False.
+            # Even if some backends may return a scalar view of the original, we chose to be
+            # strict here beceause some other backends, such as numpy, definitely don't.
+            tup_idx = self._idx if isinstance(self._idx, tuple) else (self._idx,)
+            if any(
+                i is not None and i is not Ellipsis and not isinstance(i, slice)
+                for i in tup_idx
+            ):
+                msg = "get() with a scalar index typically returns a copy"
+                raise ValueError(msg)
+
+            # Note: this is not the same list of backends as is_writeable_array()
+            if is_dask_array(x) or is_jax_array(x) or is_pydata_sparse_array(x):
+                msg = f"get() on {array_namespace(x)} arrays always returns a copy"
+                raise ValueError(msg)
+
+        if is_jax_array(x):
+            # Use JAX's at[] or other library that with the same duck-type API
+            return x.at[self._idx].get()
+
+        if xp is None:
+            xp = array_namespace(x)
+        # Note: when idx is a boolean mask, numpy always returns a deep copy.
+        # However, some backends may legitimately return a view when the mask can
+        # be downgraded to a slice, e.g. a[[True, True, False]] -> a[:2].
+        # Err on the side of caution and perform a double-copy in numpy.
+        return xp.asarray(x[self._idx], copy=copy)
+
     def _update_common(
         self,
         at_op: str,
@@ -701,23 +771,7 @@ def _update_common(
         If the operation can be resolved by at[], (return value, None)
         Otherwise, (None, preprocessed x)
         """
-        x, idx = self._x, self._idx
-
-        if idx is _undef:
-            msg = (
-                "Index has not been set.\n"
-                "Usage: either\n"
-                "    at(x, idx).set(value)\n"
-                "or\n"
-                "    at(x)[idx].set(value)\n"
-                "(same for all other methods)."
-            )
-            raise TypeError(msg)
-
-        if copy not in (True, False, None):
-            msg = f"copy must be True, False, or None; got {copy!r}"  # pyright: ignore[reportUnreachable]
-            raise ValueError(msg)
-
+        x = self._x
         if copy is None:
             writeable = is_writeable_array(x)
             copy = not writeable
@@ -758,6 +812,7 @@ def set(
         xp: ModuleType | None = None,
     ) -> Array:
         """Apply ``x[idx] = y`` and return the update array"""
+        self._check_args(copy=copy)
         res, x = self._update_common("set", y, copy=copy, xp=xp)
         if res is not None:
             return res
@@ -785,6 +840,7 @@ def _iop(
         Consider for example when x is a numpy array and idx is a fancy index, which
         triggers a deep copy on __getitem__.
         """
+        self._check_args(copy=copy)
         res, x = self._update_common(at_op, y, copy=copy, xp=xp)
         if res is not None:
             return res

src/array_api_extra/_lib/_compat.py

@@ -6,20 +6,29 @@
     from ..._array_api_compat_vendor import (  # pyright: ignore[reportMissingImports]
         array_namespace,
         device,
+        is_array_api_obj,
+        is_dask_array,
         is_jax_array,
+        is_pydata_sparse_array,
         is_writeable_array,
     )
 except ImportError:
     from array_api_compat import (  # pyright: ignore[reportMissingTypeStubs]
         array_namespace,
         device,
+        is_array_api_obj,
+        is_dask_array,
         is_jax_array,
+        is_pydata_sparse_array,
         is_writeable_array,
     )
 
 __all__ = (
     "array_namespace",
     "device",
+    "is_array_api_obj",
+    "is_dask_array",
     "is_jax_array",
+    "is_pydata_sparse_array",
     "is_writeable_array",
 )

src/array_api_extra/_lib/_compat.pyi

@@ -11,5 +11,8 @@ def array_namespace(
     use_compat: bool | None = None,
 ) -> ArrayModule: ...
 def device(x: Array, /) -> Device: ...
+def is_array_api_obj(x: object, /) -> bool: ...
+def is_dask_array(x: object, /) -> bool: ...
 def is_jax_array(x: object, /) -> bool: ...
+def is_pydata_sparse_array(x: object, /) -> bool: ...
 def is_writeable_array(x: object, /) -> bool: ...

tests/test_at.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from contextlib import contextmanager
+from contextlib import contextmanager, suppress
 from importlib import import_module
 from typing import TYPE_CHECKING, Final
 
@@ -9,6 +9,7 @@
 from array_api_compat import (  # type: ignore[import-untyped]  # pyright: ignore[reportMissingTypeStubs]
     array_namespace,
     is_dask_array,
+    is_numpy_array,
     is_pydata_sparse_array,
     is_writeable_array,
 )
@@ -98,6 +99,71 @@ def test_update_ops(
         assert_array_equal(y, expect)
 
 
+@pytest.mark.parametrize("copy", [True, False, None])
+def test_get(array: Array, copy: bool | None):
+    expect_copy = copy
+
+    # dask is mutable, but __getitem__ never returns a view
+    if is_dask_array(array):
+        if copy is False:
+            with pytest.raises(ValueError, match="always returns a copy"):
+                at(array, slice(2)).get(copy=False)
+            return
+        expect_copy = True
+
+    # get(copy=False) on a read-only numpy array returns a read-only view
+    if is_numpy_array(array) and not copy and not array.flags.writeable:
+        out = at(array, slice(2)).get(copy=copy)
+        assert_array_equal(out, [10.0, 20.0])
+        assert out.base is array
+        assert not out.flags.writeable
+        return
+
+    with assert_copy(array, expect_copy):
+        y = at(array, slice(2)).get(copy=copy)
+        assert isinstance(y, type(array))
+        assert_array_equal(y, [10.0, 20.0])
+        # Let assert_copy test that y is a view or copy
+        with suppress(TypeError, ValueError):
+            y[:] = 40
+
+
+def test_get_scalar_nocopy(array: Array):
+    """get(copy=False) with a scalar index always raises, because some backends
+    such as numpy and sparse return a np.generic instead of a scalar view
+    """
+    with pytest.raises(ValueError, match="scalar"):
+        at(array)[0].get(copy=False)
+    with pytest.raises(ValueError, match="scalar"):
+        at(array)[(0,)].get(copy=False)
+    with pytest.raises(ValueError, match="scalar"):
+        at(array)[..., 0].get(copy=False)
+
+
+def test_get_bool_indices(array: Array):
+    """get() with a boolean array index always returns a copy"""
+    # sparse violates the array API as it doesn't support
+    # a boolean index that is another sparse array.
+    # dask with dask index has NaN size, which complicates testing.
+    if is_pydata_sparse_array(array) or is_dask_array(array):
+        xp = np
+    else:
+        xp = array_namespace(array)
+    idx = xp.asarray([True, False, True])
+
+    with pytest.raises(ValueError, match="copy"):
+        at(array, idx).get(copy=False)
+
+    assert_array_equal(at(array, idx).get(), [10.0, 30.0])
+
+    with assert_copy(array, True):
+        y = at(array, idx).get(copy=True)
+        assert_array_equal(y, [10.0, 30.0])
+        # Let assert_copy test that y is a view or copy
+        with suppress(TypeError, ValueError):
+            y[:] = 40
+
+
 def test_copy_invalid():
     a = np.asarray([1, 2, 3])
     with pytest.raises(ValueError, match="copy"):
@@ -106,6 +172,7 @@ def test_copy_invalid():
 
 def test_xp():
     a = np.asarray([1, 2, 3])
+    at(a, 0).get(xp=np)
     at(a, 0).set(4, xp=np)
     at(a, 0).add(4, xp=np)
     at(a, 0).subtract(4, xp=np)

vendor_tests/test_vendor.py

@@ -8,14 +8,21 @@ def test_vendor_compat():
     from ._array_api_compat_vendor import (  # type: ignore[attr-defined]
         array_namespace,
         device,
+        is_array_api_obj,
+        is_dask_array,
         is_jax_array,
+        is_pydata_sparse_array,
         is_writeable_array,
     )
 
     x = xp.asarray([1, 2, 3])
     assert array_namespace(x) is xp
     device(x)
+    assert is_array_api_obj(x)
+    assert not is_array_api_obj(123)
+    assert not is_dask_array(x)
     assert not is_jax_array(x)
+    assert not is_pydata_sparse_array(x)
     assert is_writeable_array(x)