Factor out the list of wrapped libraries in the tests

This does remove jax testing from the isdtype test, since we should not test
jax functionality for non-helper functions.

Author: asmeurer

tests/_helpers.py

@@ -5,6 +5,9 @@
 import pytest
 
 
+wrapped_libraries = ["numpy", "cupy", "torch", "dask.array"]
+all_libraries = wrapped_libraries + "jax.numpy"
+
 def import_(library, wrapper=False):
     if 'jax' in library and sys.version_info < (3, 9):
         pytest.skip('JAX array API support does not support Python 3.8')

tests/test_all.py

@@ -12,11 +12,11 @@
 
 import sys
 
-from ._helpers import import_
+from ._helpers import import_, wrapped_libraries
 
 import pytest
 
-@pytest.mark.parametrize("library", ["common", "cupy", "numpy", "torch", "dask.array"])
+@pytest.mark.parametrize("library", ["common"] + wrapped_libraries)
 def test_all(library):
     import_(library, wrapper=True)
 

tests/test_array_namespace.py

@@ -9,9 +9,9 @@
 import array_api_compat
 from array_api_compat import array_namespace
 
-from ._helpers import import_
+from ._helpers import import_, all_libraries
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+@pytest.mark.parametrize("library", all_libraries)
 @pytest.mark.parametrize("api_version", [None, "2021.12"])
 def test_array_namespace(library, api_version):
     xp = import_(library)

tests/test_common.py

@@ -3,7 +3,7 @@
 
 from array_api_compat import is_array_api_obj, device, to_device
 
-from ._helpers import import_
+from ._helpers import import_, wrapped_libraries, all_libraries
 
 import pytest
 import numpy as np
@@ -29,7 +29,7 @@ def test_is_xp_array(library, func):
 
     assert is_array_api_obj(x)
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+@pytest.mark.parametrize("library", all_libraries)
 def test_device(library):
     xp = import_(library, wrapper=True)
 
@@ -43,7 +43,7 @@ def test_device(library):
     assert device(x) == device(x2)
 
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array"])
+@pytest.mark.parametrize("library", wrapped_libraries)
 def test_to_device_host(library):
     # different libraries have different semantics
     # for DtoH transfers; ensure that we support a portable

tests/test_isdtype.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from ._helpers import import_
+from ._helpers import import_, wrapped_libraries
 
 # Check the known dtypes by their string names
 
@@ -64,7 +64,7 @@ def isdtype_(dtype_, kind):
     assert type(res) is bool #  noqa: E721
     return res
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+@pytest.mark.parametrize("library", wrapped_libraries)
 def test_isdtype_spec_dtypes(library):
     xp = import_(library, wrapper=True)
 
@@ -98,7 +98,7 @@ def test_isdtype_spec_dtypes(library):
     'bfloat16',
 ]
 
-@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+@pytest.mark.parametrize("library", wrapped_libraries)
 @pytest.mark.parametrize("dtype_", additional_dtypes)
 def test_isdtype_additional_dtypes(library, dtype_):
     xp = import_(library, wrapper=True)