Allow iteration on 1-D arrays

upstream code (see scipy/scipy#21074). The standard does not define iteration,
but it does define __getitem__, and the default Python __iter__ implements
iteration when getitem is defined as a[0], a[1], ..., implying that iteration
ought to work for 1-D arrays. Iteration is still disallowed for higher
dimensional arrays, since getitem would not necessarily work with a single
integer index (and this is the case that is controversial). In those cases,
the new unstack() function would be preferable.

At best it would be good to get upstream clarification from the standard
whether iteration should always work or not before disallowing 1-D array
iteration.

Author: asmeurer

array_api_strict/_array_object.py

@@ -677,10 +677,16 @@ def __iter__(self: Array, /):
         """
         Performs the operation __iter__.
         """
-        # Manually disable iteration, since __getitem__ raises IndexError on
-        # things like ones((3, 3))[0], which causes list(ones((3, 3))) to give
-        # [].
-        raise TypeError("array iteration is not allowed in array-api-strict")
+        # Manually disable iteration on higher dimensional arrays, since
+        # __getitem__ raises IndexError on things like ones((3, 3))[0], which
+        # causes list(ones((3, 3))) to give [].
+        if self.ndim > 1:
+            raise TypeError("array iteration is not allowed in array-api-strict")
+        # Allow iteration for 1-D arrays. The array API doesn't strictly
+        # define __iter__, but it doesn't disallow it. The default Python
+        # behavior is to implement iter as a[0], a[1], ... when __getitem__ is
+        # implemented, which implies iteration on 1-D arrays.
+        return (Array._new(i) for i in self._array)
 
     def __le__(self: Array, other: Union[int, float, Array], /) -> Array:
         """

array_api_strict/tests/test_array_object.py

@@ -1,4 +1,5 @@
 import operator
+from builtins import all as all_
 
 from numpy.testing import assert_raises, suppress_warnings
 import numpy as np
@@ -21,6 +22,7 @@
     int32,
     int64,
     uint64,
+    float64,
     bool as bool_,
 )
 from .._flags import set_array_api_strict_flags
@@ -423,8 +425,12 @@ def test_array_namespace():
     pytest.raises(ValueError, lambda: a.__array_namespace__(api_version="2021.11"))
     pytest.raises(ValueError, lambda: a.__array_namespace__(api_version="2024.12"))
 
-def test_no_iter():
-    pytest.raises(TypeError, lambda: iter(ones(3)))
+def test_iter():
+    pytest.raises(TypeError, lambda: iter(asarray(3)))
+    assert list(ones(3)) == [asarray(1.), asarray(1.), asarray(1.)]
+    assert all_(isinstance(a, Array) for a in iter(ones(3)))
+    assert all_(a.shape == () for a in iter(ones(3)))
+    assert all_(a.dtype == float64 for a in iter(ones(3)))
     pytest.raises(TypeError, lambda: iter(ones((3, 3))))
 
 @pytest.mark.parametrize("api_version", ['2021.12', '2022.12', '2023.12'])