test_cupy: test when CuPy isn't available

Author: madsbk

zarr/tests/test_cupy.py

@@ -13,9 +13,6 @@
 from zarr.storage import DirectoryStore, MemoryStore, Store, ZipStore
 
 
-cupy = pytest.importorskip("cupy")
-
-
 class CuPyCPUCompressor(Codec):
     """CPU compressor for CuPy arrays
 
@@ -71,6 +68,28 @@ def from_config(cls, config):
 register_codec(CuPyCPUCompressor)
 
 
+class MyArray(np.ndarray):
+    """Dummy array class to test the `meta_array` argument
+
+    Useful when CuPy isn't available.
+
+    This class also makes some of the functions from the numpy
+    module available.
+    """
+
+    testing = np.testing
+
+    @classmethod
+    def arange(cls, size):
+        ret = cls(shape=(size,), dtype="int64")
+        ret[:] = range(size)
+        return ret
+
+    @classmethod
+    def empty(cls, shape):
+        return cls(shape=shape)
+
+
 def init_compressor(compressor) -> CuPyCPUCompressor:
     if compressor:
         compressor = getattr(zarr.codecs, compressor)()
@@ -85,93 +104,113 @@ def init_store(tmp_path, store_type) -> Optional[Store]:
     return None
 
 
-@pytest.mark.parametrize("compressor", [None, "Zlib", "Blosc"])
+def ensure_cls(obj):
+    if isinstance(obj, str):
+        module, cls_name = obj.rsplit(".", maxsplit=1)
+        return getattr(pytest.importorskip(module), cls_name)
+    return obj
+
+
+def ensure_module(module):
+    if isinstance(module, str):
+        return pytest.importorskip(module)
+    return module
+
+
+param_module_and_compressor = [
+    (MyArray, None),
+    ("cupy", init_compressor(None)),
+    ("cupy", init_compressor("Zlib")),
+    ("cupy", init_compressor("Blosc")),
+]
+
+
+@pytest.mark.parametrize("module, compressor", param_module_and_compressor)
 @pytest.mark.parametrize("store_type", [None, DirectoryStore, MemoryStore, ZipStore])
-def test_array(tmp_path, compressor, store_type):
-    compressor = init_compressor(compressor)
+def test_array(tmp_path, module, compressor, store_type):
+    xp = ensure_module(module)
 
-    # with cupy array
     store = init_store(tmp_path / "from_cupy_array", store_type)
-    a = cupy.arange(100)
-    z = array(
-        a, chunks=10, compressor=compressor, store=store, meta_array=cupy.empty(())
-    )
+    a = xp.arange(100)
+    z = array(a, chunks=10, compressor=compressor, store=store, meta_array=xp.empty(()))
     assert a.shape == z.shape
     assert a.dtype == z.dtype
     assert isinstance(a, type(z[:]))
-    assert isinstance(z.meta_array, type(cupy.empty(())))
-    cupy.testing.assert_array_equal(a, z[:])
+    assert isinstance(z.meta_array, type(xp.empty(())))
+    xp.testing.assert_array_equal(a, z[:])
 
     # with array-like
     store = init_store(tmp_path / "from_list", store_type)
     a = list(range(100))
-    z = array(
-        a, chunks=10, compressor=compressor, store=store, meta_array=cupy.empty(())
-    )
+    z = array(a, chunks=10, compressor=compressor, store=store, meta_array=xp.empty(()))
     assert (100,) == z.shape
     assert np.asarray(a).dtype == z.dtype
-    cupy.testing.assert_array_equal(a, z[:])
+    xp.testing.assert_array_equal(a, z[:])
 
     # with another zarr array
     store = init_store(tmp_path / "from_another_store", store_type)
-    z2 = array(z, compressor=compressor, store=store, meta_array=cupy.empty(()))
+    z2 = array(z, compressor=compressor, store=store, meta_array=xp.empty(()))
     assert z.shape == z2.shape
     assert z.chunks == z2.chunks
     assert z.dtype == z2.dtype
-    cupy.testing.assert_array_equal(z[:], z2[:])
+    xp.testing.assert_array_equal(z[:], z2[:])
 
 
-@pytest.mark.parametrize("compressor", [None, "Zlib", "Blosc"])
-def test_empty(compressor):
+@pytest.mark.parametrize("module, compressor", param_module_and_compressor)
+def test_empty(module, compressor):
+    xp = ensure_module(module)
     z = empty(
         100,
         chunks=10,
         compressor=init_compressor(compressor),
-        meta_array=cupy.empty(()),
+        meta_array=xp.empty(()),
     )
     assert (100,) == z.shape
     assert (10,) == z.chunks
 
 
-@pytest.mark.parametrize("compressor", [None, "Zlib", "Blosc"])
-def test_zeros(compressor):
+@pytest.mark.parametrize("module, compressor", param_module_and_compressor)
+def test_zeros(module, compressor):
+    xp = ensure_module(module)
     z = zeros(
         100,
         chunks=10,
         compressor=init_compressor(compressor),
-        meta_array=cupy.empty(()),
+        meta_array=xp.empty(()),
     )
     assert (100,) == z.shape
     assert (10,) == z.chunks
-    cupy.testing.assert_array_equal(np.zeros(100), z[:])
+    xp.testing.assert_array_equal(np.zeros(100), z[:])
 
 
-@pytest.mark.parametrize("compressor", [None, "Zlib", "Blosc"])
-def test_ones(compressor):
+@pytest.mark.parametrize("module, compressor", param_module_and_compressor)
+def test_ones(module, compressor):
+    xp = ensure_module(module)
     z = ones(
         100,
         chunks=10,
         compressor=init_compressor(compressor),
-        meta_array=cupy.empty(()),
+        meta_array=xp.empty(()),
     )
     assert (100,) == z.shape
     assert (10,) == z.chunks
-    cupy.testing.assert_array_equal(np.ones(100), z[:])
+    xp.testing.assert_array_equal(np.ones(100), z[:])
 
 
-@pytest.mark.parametrize("compressor", [None, "Zlib", "Blosc"])
-def test_full(compressor):
+@pytest.mark.parametrize("module, compressor", param_module_and_compressor)
+def test_full(module, compressor):
+    xp = ensure_module(module)
     z = full(
         100,
         chunks=10,
         fill_value=42,
         dtype="i4",
         compressor=init_compressor(compressor),
-        meta_array=cupy.empty(()),
+        meta_array=xp.empty(()),
     )
     assert (100,) == z.shape
     assert (10,) == z.chunks
-    cupy.testing.assert_array_equal(np.full(100, fill_value=42, dtype="i4"), z[:])
+    xp.testing.assert_array_equal(np.full(100, fill_value=42, dtype="i4"), z[:])
 
     # nan
     z = full(
@@ -180,21 +219,22 @@ def test_full(compressor):
         fill_value=np.nan,
         dtype="f8",
         compressor=init_compressor(compressor),
-        meta_array=cupy.empty(()),
+        meta_array=xp.empty(()),
     )
     assert np.all(np.isnan(z[:]))
 
 
-@pytest.mark.parametrize("compressor", [None, "Zlib", "Blosc"])
+@pytest.mark.parametrize("module, compressor", param_module_and_compressor)
 @pytest.mark.parametrize("store_type", [None, DirectoryStore, MemoryStore, ZipStore])
-def test_group(tmp_path, compressor, store_type):
+def test_group(tmp_path, module, compressor, store_type):
+    xp = ensure_module(module)
     store = init_store(tmp_path, store_type)
-    g = open_group(store, meta_array=cupy.empty(()))
+    g = open_group(store, meta_array=xp.empty(()))
     g.ones("data", shape=(10, 11), dtype=int, compressor=init_compressor(compressor))
     a = g["data"]
     assert a.shape == (10, 11)
     assert a.dtype == int
     assert isinstance(a, Array)
-    assert isinstance(a[:], cupy.ndarray)
+    assert isinstance(a[:], xp.ndarray)
     assert (a[:] == 1).all()
-    assert isinstance(g.meta_array, type(cupy.empty(())))
+    assert isinstance(g.meta_array, type(xp.empty(())))