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