Temporarily enable __array__ in asarray so that parsing list of lists of Array can work

Author: asmeurer

array_api_strict/_array_object.py

@@ -53,6 +53,8 @@ def __repr__(self):
 
 _default = object()
 
+_allow_array = False
+
 class Array:
     """
     n-d array object for the array API namespace.
@@ -135,6 +137,22 @@ def __repr__(self: Array, /) -> str:
     # lead to code assuming np.asarray(other_array) would always work in the
     # standard.
     def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None) -> npt.NDArray[Any]:
+        # We have to allow this to be internally enabled as there's no other
+        # easy way to parse a list of Array objects in asarray().
+        if _allow_array:
+            # copy keyword is new in 2.0.0; for older versions don't use it
+            # retry without that keyword.
+            if np.__version__[0] < '2':
+                return np.asarray(self._array, dtype=dtype)
+            elif np.__version__.startswith('2.0.0-dev0'):
+                # Handle dev version for which we can't know based on version
+                # number whether or not the copy keyword is supported.
+                try:
+                    return np.asarray(self._array, dtype=dtype, copy=copy)
+                except TypeError:
+                    return np.asarray(self._array, dtype=dtype)
+            else:
+                return np.asarray(self._array, dtype=dtype, copy=copy)
         raise ValueError("Conversion from an array_api_strict array to a NumPy ndarray is not supported")
 
     # These are various helper functions to make the array behavior match the

array_api_strict/_creation_functions.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-
+from contextlib import contextmanager
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 if TYPE_CHECKING:
@@ -16,6 +16,19 @@
 
 import numpy as np
 
+@contextmanager
+def allow_array():
+    """
+    Temporarily enable Array.__array__. This is needed for np.array to parse
+    list of lists of Array objects.
+    """
+    from . import _array_object
+    original_value = _array_object._allow_array
+    try:
+        _array_object._allow_array = True
+        yield
+    finally:
+        _array_object._allow_array = original_value
 
 def _check_valid_dtype(dtype):
     # Note: Only spelling dtypes as the dtype objects is supported.
@@ -94,7 +107,8 @@ def asarray(
         # Give a better error message in this case. NumPy would convert this
         # to an object array. TODO: This won't handle large integers in lists.
         raise OverflowError("Integer out of bounds for array dtypes")
-    res = np.array(obj, dtype=_np_dtype, copy=copy)
+    with allow_array():
+        res = np.array(obj, dtype=_np_dtype, copy=copy)
     return Array._new(res)
 
 

array_api_strict/tests/test_creation_functions.py

@@ -22,7 +22,7 @@
     zeros,
     zeros_like,
 )
-from .._dtypes import float32, float64
+from .._dtypes import int16, float32, float64
 from .._array_object import Array, CPU_DEVICE
 from .._flags import set_array_api_strict_flags
 
@@ -97,6 +97,19 @@ def test_asarray_copy():
     a[0] = 0
     assert all(b[0] == 0)
 
+def test_asarray_list_of_lists():
+    a = asarray(1, dtype=int16)
+    b = asarray([1], dtype=int16)
+    res = asarray([a, a])
+    assert res.shape == (2,)
+    assert res.dtype == int16
+    assert all(res == asarray([1, 1]))
+
+    res = asarray([b, b])
+    assert res.shape == (2, 1)
+    assert res.dtype == int16
+    assert all(res == asarray([[1], [1]]))
+
 def test_arange_errors():
     arange(1, device=CPU_DEVICE)  # Doesn't error
     assert_raises(ValueError, lambda: arange(1, device="cpu"))