diff --git a/cubed_xarray/cubedmanager.py b/cubed_xarray/cubedmanager.py index 0292aa3..d81e32f 100644 --- a/cubed_xarray/cubedmanager.py +++ b/cubed_xarray/cubedmanager.py @@ -210,23 +210,12 @@ def store( 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, - **kwargs, + **kwargs, # regions passed through here ) diff --git a/cubed_xarray/tests/test_wrapping.py b/cubed_xarray/tests/test_wrapping.py index 0abee8f..78ee5e8 100644 --- a/cubed_xarray/tests/test_wrapping.py +++ b/cubed_xarray/tests/test_wrapping.py @@ -1,6 +1,7 @@ import sys import cubed +import numpy as np import pytest import xarray as xr from cubed.runtime.create import create_executor @@ -75,6 +76,29 @@ def test_to_zarr(tmpdir, executor): assert_allclose(original, computed) +# based on test_write_region +def test_write_region(tmpdir): + zeros = xr.Dataset({"u": (("x",), np.zeros(10))}).chunk( + 2, chunked_array_type="cubed" + ) + nonzeros = xr.Dataset({"u": (("x",), np.arange(1, 11))}).chunk( + 2, chunked_array_type="cubed" + ) + + store = tmpdir / "out.zarr" + zeros.to_zarr( + store, + encoding={"u": dict(chunks=2)}, + ) + with xr.open_zarr(store) as actual: + assert_identical(actual, zeros) + for i in range(0, 10, 2): + region = {"x": slice(i, i + 2)} + nonzeros.isel(region).to_zarr(store, region=region) + with xr.open_zarr(store) as actual: + assert_identical(actual, nonzeros) + + def test_dataset_accessor_visualize(tmp_path): spec = cubed.Spec(allowed_mem="200MB")