From d944059a6e43add1d32183c1611e830eb51eed64 Mon Sep 17 00:00:00 2001
From: Tom White <tom.e.white@gmail.com>
Date: Mon, 9 Mar 2026 09:29:20 +0000
Subject: [PATCH] Allow `to_zarr` to be a lazy operation by passing
 `compute=False`

---
 cubed/core/ops.py        | 11 +++++++++--
 cubed/tests/test_core.py | 19 +++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index 186afa70..ec07fd4f 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -277,7 +277,9 @@ def __iter__(self):
         return out
 
 
-def to_zarr(x: "Array", store, path=None, region=None, executor=None, **kwargs):
+def to_zarr(
+    x: "Array", store, path=None, region=None, compute=True, *, executor=None, **kwargs
+):
     """Save an array to Zarr storage.
 
     Note that this operation is eager, and will run the computation
@@ -293,12 +295,17 @@ def to_zarr(x: "Array", store, path=None, region=None, executor=None, **kwargs):
         Group path
     region : tuple of slices, optional
         The region of data that should be written to in target.
+    compute : boolean, optional
+        If True compute immediately, return array otherwise.
     executor : cubed.runtime.types.Executor, optional
         The executor to use to run the computation.
         Defaults to using the in-process Python executor.
     """
     out = _store_array(x, store, path=path, region=region)
-    out.compute(executor=executor, _return_in_memory_array=False, **kwargs)
+    if compute:
+        out.compute(executor=executor, _return_in_memory_array=False, **kwargs)
+    else:
+        return out
 
 
 def blockwise(
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index e00fcffa..a81fbb5c 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -224,6 +224,25 @@ def test_to_zarr_array(tmp_path, spec, executor):
     assert_array_equal(res[:], np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
 
 
+def test_to_zarr_lazy_compute(tmp_path, spec):
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    store = tmp_path / "output.zarr"
+    z = create_zarr(
+        np.zeros(a.shape, dtype=a.dtype),
+        chunks=(2, 2),
+        store=store,
+    )
+    b = cubed.to_zarr(a, z, compute=False)
+    # target has not been computed yet
+    res = open_storage_array(store, mode="r")
+    with pytest.raises(AssertionError):
+        assert_array_equal(res[:], np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
+
+    b.compute()
+    res = open_storage_array(store, mode="r")
+    assert_array_equal(res[:], np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
+
+
 def test_to_zarr_region(tmp_path, spec, executor):
     a = xp.asarray(
         [