Reuse alternate index syntax

data-apis · Dec 18, 2024 · d0bbc88 · d0bbc88
1 parent c9d45af
commit d0bbc88
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/src/array_api_extra/_funcs.py b/src/array_api_extra/_funcs.py
@@ -672,8 +672,7 @@ def __getitem__(self, idx: Index, /) -> "at":  # numpydoc ignore=PR01,RT01
         if self._idx is not _undef:
             msg = "Index has already been set"
             raise ValueError(msg)
-        self._idx = idx
-        return self
+        return at(self._x, idx)
 
     def _update_common(
         self,

diff --git a/tests/test_at.py b/tests/test_at.py
@@ -127,6 +127,11 @@ def test_alternate_index_syntax():
     a = np.asarray([1, 2, 3])
     assert_array_equal(at(a, 0).set(4), [4, 2, 3])
     assert_array_equal(at(a)[0].set(4), [4, 2, 3])
+
+    a_at = at(a)
+    assert_array_equal(a_at[0].add(1), [2, 2, 3])
+    assert_array_equal(a_at[1].add(2), [1, 4, 3])
+
     with pytest.raises(ValueError, match="Index"):
         at(a).set(4)
     with pytest.raises(ValueError, match="Index"):