Remove get()

Author: crusaderky

src/array_api_extra/_funcs.py

@@ -11,10 +11,7 @@
 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,
 )
 
@@ -689,73 +686,6 @@ 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,
@@ -771,7 +701,23 @@ def _update_common(
         If the operation can be resolved by at[], (return value, None)
         Otherwise, (None, preprocessed x)
         """
-        x = self._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)
+
         if copy is None:
             writeable = is_writeable_array(x)
             copy = not writeable
@@ -812,7 +758,6 @@ 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
@@ -840,7 +785,6 @@ 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,29 +6,20 @@
     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,8 +11,5 @@ 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, suppress
+from contextlib import contextmanager
 from importlib import import_module
 from typing import TYPE_CHECKING, Final
 
@@ -9,7 +9,6 @@
 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,
 )
@@ -99,71 +98,6 @@ 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"):
@@ -172,7 +106,6 @@ 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,21 +8,14 @@ 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)