Rename import_or_skip_cupy to import_ and move jax logic into it

Author: asmeurer

tests/_helpers.py

@@ -3,7 +3,14 @@
 import pytest
 
 
-def import_or_skip_cupy(library):
-    if "cupy" in library:
+def import_(library, wrapper=False):
+    if library == 'cupy':
         return pytest.importorskip(library)
+
+    if wrapper:
+        if 'jax' in library:
+            library = 'jax.experimental.array_api'
+        else:
+            library = 'array_api_compat.' + library
+
     return import_module(library)

tests/test_array_namespace.py

@@ -5,13 +5,13 @@
 import array_api_compat
 from array_api_compat import array_namespace
 
-from ._helpers import import_or_skip_cupy
+from ._helpers import import_
 
 
 @pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array"])
 @pytest.mark.parametrize("api_version", [None, "2021.12"])
 def test_array_namespace(library, api_version):
-    xp = import_or_skip_cupy(library)
+    xp = import_(library)
 
     array = xp.asarray([1.0, 2.0, 3.0])
     namespace = array_api_compat.array_namespace(array, api_version=api_version)

tests/test_helpers.py

@@ -2,7 +2,7 @@
                               is_dask_array, is_jax_array, is_array_api_obj,
                               device, to_device)
 
-from ._helpers import import_or_skip_cupy
+from ._helpers import import_
 
 import pytest
 import numpy as np
@@ -19,7 +19,7 @@
 @pytest.mark.parametrize('library', is_functions.keys())
 @pytest.mark.parametrize('func', is_functions.values())
 def test_is_xp_array(library, func):
-    lib = import_or_skip_cupy(library)
+    lib = import_(library)
     is_func = globals()[func]
 
     x = lib.asarray([1, 2, 3])
@@ -33,10 +33,7 @@ def test_device(library):
     if library == "dask.array":
         pytest.xfail("device() needs to be fixed for dask")
 
-    if library == "jax.numpy":
-        xp = import_or_skip_cupy('jax.experimental.array_api')
-    else:
-        xp = import_or_skip_cupy('array_api_compat.' + library)
+    xp = import_(library, wrapper=True)
 
     # We can't test much for device() and to_device() other than that
     # x.to_device(x.device) works.
@@ -54,7 +51,7 @@ def test_to_device_host(library):
     # for DtoH transfers; ensure that we support a portable
     # shim for common array libs
     # see: https://github.com/scipy/scipy/issues/18286#issuecomment-1527552919
-    xp = import_or_skip_cupy('array_api_compat.' + library)
+    xp = import_(library, wrapper=True)
 
     expected = np.array([1, 2, 3])
     x = xp.asarray([1, 2, 3])

tests/test_isdtype.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from ._helpers import import_or_skip_cupy
+from ._helpers import import_
 
 # Check the known dtypes by their string names
 
@@ -66,10 +66,7 @@ def isdtype_(dtype_, kind):
 
 @pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
 def test_isdtype_spec_dtypes(library):
-    if library == "jax.numpy":
-        xp = import_or_skip_cupy('jax.experimental.array_api')
-    else:
-        xp = import_or_skip_cupy('array_api_compat.' + library)
+    xp = import_(library, wrapper=True)
 
     isdtype = xp.isdtype
 
@@ -104,10 +101,7 @@ def test_isdtype_spec_dtypes(library):
 @pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
 @pytest.mark.parametrize("dtype_", additional_dtypes)
 def test_isdtype_additional_dtypes(library, dtype_):
-    if library == "jax.numpy":
-        xp = import_or_skip_cupy('jax.experimental.array_api')
-    else:
-        xp = import_or_skip_cupy('array_api_compat.' + library)
+    xp = import_(library, wrapper=True)
 
     isdtype = xp.isdtype