Handle kwargs better in store (#14)

* Handle kwargs better in store

* Add a test

* Only run with processes executor on Python 3.11

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: tomwhite

cubed_xarray/__init__.py

@@ -1,6 +1,5 @@
 from importlib.metadata import version
 
-
 try:
     __version__ = version("cubed-xarray")
 except Exception:

cubed_xarray/cubedmanager.py

@@ -1,18 +1,15 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Callable, Union
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Union
 
 import numpy as np
-
 from tlz import partition
-
 from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint
 
-
 if TYPE_CHECKING:
-    from xarray.core.types import T_Chunks, T_NormalizedChunks
     from cubed import Array as CubedArray
+    from xarray.core.types import T_Chunks, T_NormalizedChunks
 
 
 class CubedManager(ChunkManagerEntrypoint["CubedArray"]):
@@ -204,6 +201,27 @@ def store(
         """Used when writing to any backend."""
         from cubed.core.ops import store
 
+        compute = kwargs.pop("compute", True)
+        if not compute:
+            raise NotImplementedError("Delayed compute is not supported.")
+
+        lock = kwargs.pop("lock", None)
+        if lock:
+            raise NotImplementedError("Locking is not supported.")
+
+        regions = kwargs.pop("regions", None)
+        if regions:
+            # regions is either a tuple of slices or a collection of tuples of slices
+            if isinstance(regions, tuple):
+                regions = [regions]
+            for t in regions:
+                if not all(r == slice(None) for r in t):
+                    raise NotImplementedError(
+                        "Only whole slices are supported for regions."
+                    )
+
+        kwargs.pop("flush", None)  # not used
+
         return store(
             sources,
             targets,

cubed_xarray/tests/test_wrapping.py

@@ -1,22 +1,63 @@
+import sys
+
+import cubed
+import pytest
 import xarray as xr
+from cubed.runtime.create import create_executor
 from xarray.namedarray.parallelcompat import list_chunkmanagers
-import cubed
+from xarray.tests import assert_allclose, create_test_data
 
 from cubed_xarray.cubedmanager import CubedManager
 
+EXECUTORS = [create_executor("single-threaded")]
+
+if sys.version_info >= (3, 11):
+    EXECUTORS.append(create_executor("processes"))
+
+
+@pytest.fixture(
+    scope="module",
+    params=EXECUTORS,
+    ids=[executor.name for executor in EXECUTORS],
+)
+def executor(request):
+    return request.param
+
 
 class TestDiscoverCubedManager:
     def test_list_cubedmanager(self):
         chunkmanagers = list_chunkmanagers()
-        assert 'cubed' in chunkmanagers
-        assert isinstance(chunkmanagers['cubed'], CubedManager)
+        assert "cubed" in chunkmanagers
+        assert isinstance(chunkmanagers["cubed"], CubedManager)
 
     def test_chunk(self):
-        da = xr.DataArray([1, 2], dims='x')
-        chunked = da.chunk(x=1, chunked_array_type='cubed')
+        da = xr.DataArray([1, 2], dims="x")
+        chunked = da.chunk(x=1, chunked_array_type="cubed")
         assert isinstance(chunked.data, cubed.Array)
-        assert chunked.chunksizes == {'x': (1, 1)}
+        assert chunked.chunksizes == {"x": (1, 1)}
 
     # TODO test cubed is default when dask not installed
 
     # TODO test dask is default over cubed when both installed
+
+
+def test_to_zarr(tmpdir, executor):
+    spec = cubed.Spec(allowed_mem="200MB", executor=executor)
+
+    original = create_test_data().chunk(
+        chunked_array_type="cubed", from_array_kwargs={"spec": spec}
+    )
+
+    filename = tmpdir / "out.zarr"
+    original.to_zarr(filename)
+
+    with xr.open_dataset(
+        filename,
+        chunks="auto",
+        engine="zarr",
+        chunked_array_type="cubed",
+        from_array_kwargs={"spec": spec},
+    ) as restored:
+        assert isinstance(restored.var1.data, cubed.Array)
+        computed = restored.compute()
+        assert_allclose(original, computed)