TST: Test size()

Author: crusaderky

tests/_helpers.py

@@ -11,7 +11,7 @@
     all_libraries.append('sparse')
 
 def import_(library, wrapper=False):
-    if library == 'cupy':
+    if library in ('cupy', 'ndonnx'):
         pytest.importorskip(library)
     if wrapper:
         if 'jax' in library:

tests/test_common.py

@@ -5,8 +5,9 @@
     is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace,
 )
 
-from array_api_compat import device, is_array_api_obj, is_writeable_array, to_device
-
+from array_api_compat import (
+    device, is_array_api_obj, is_writeable_array, size, to_device
+)
 from ._helpers import import_, wrapped_libraries, all_libraries
 
 import pytest
@@ -92,6 +93,27 @@ def test_is_writeable_array_numpy():
     assert not is_writeable_array(x)
 
 
+@pytest.mark.parametrize("library", all_libraries)
+def test_size(library):
+    xp = import_(library)
+    x = xp.asarray([1, 2, 3])
+    assert size(x) == 3
+
+
+def test_size_nan():
+    xp = import_("dask.array")
+    x = xp.arange(10)
+    x = x[x < 5]
+    assert np.isnan(size(x))
+
+
+def test_size_none():
+    xp = import_("ndonnx")
+    x = xp.arange(10)
+    x = x[x < 5]
+    assert size(x) is None
+
+
 @pytest.mark.parametrize("library", all_libraries)
 def test_device(library):
     xp = import_(library, wrapper=True)