From 30468ab30faee4946018188066f1a8c01cb5b40b Mon Sep 17 00:00:00 2001 From: Tom White Date: Mon, 5 Aug 2024 17:40:34 +0100 Subject: [PATCH 1/4] Handle kwargs better in store --- cubed_xarray/cubedmanager.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/cubed_xarray/cubedmanager.py b/cubed_xarray/cubedmanager.py index 099ce89..dbb9705 100644 --- a/cubed_xarray/cubedmanager.py +++ b/cubed_xarray/cubedmanager.py @@ -204,6 +204,25 @@ 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, From 79e07b696697f79de94fcff744363d305c9b1e95 Mon Sep 17 00:00:00 2001 From: Tom White Date: Tue, 6 Aug 2024 15:45:15 +0100 Subject: [PATCH 2/4] Add a test --- cubed_xarray/tests/test_wrapping.py | 31 +++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/cubed_xarray/tests/test_wrapping.py b/cubed_xarray/tests/test_wrapping.py index d1b2a0d..202b74b 100644 --- a/cubed_xarray/tests/test_wrapping.py +++ b/cubed_xarray/tests/test_wrapping.py @@ -1,10 +1,25 @@ import xarray as xr +from xarray.tests import assert_allclose, create_test_data + from xarray.namedarray.parallelcompat import list_chunkmanagers +import pytest import cubed +from cubed.runtime.create import create_executor from cubed_xarray.cubedmanager import CubedManager +EXECUTORS = [create_executor("single-threaded"), 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() @@ -20,3 +35,19 @@ def test_chunk(self): # 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) \ No newline at end of file From 8704080bff79459f76e5834326bc8438d316e7e6 Mon Sep 17 00:00:00 2001 From: Tom White Date: Wed, 7 Aug 2024 10:50:24 +0100 Subject: [PATCH 3/4] Only run with processes executor on Python 3.11 --- cubed_xarray/cubedmanager.py | 11 ++++----- cubed_xarray/tests/test_wrapping.py | 36 ++++++++++++++++++----------- 2 files changed, 28 insertions(+), 19 deletions(-) diff --git a/cubed_xarray/cubedmanager.py b/cubed_xarray/cubedmanager.py index dbb9705..52ca284 100644 --- a/cubed_xarray/cubedmanager.py +++ b/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"]): @@ -219,7 +216,9 @@ def store( 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.") + raise NotImplementedError( + "Only whole slices are supported for regions." + ) kwargs.pop("flush", None) # not used diff --git a/cubed_xarray/tests/test_wrapping.py b/cubed_xarray/tests/test_wrapping.py index 202b74b..3124e18 100644 --- a/cubed_xarray/tests/test_wrapping.py +++ b/cubed_xarray/tests/test_wrapping.py @@ -1,15 +1,19 @@ -import xarray as xr -from xarray.tests import assert_allclose, create_test_data +import sys -from xarray.namedarray.parallelcompat import list_chunkmanagers -import pytest import cubed +import pytest +import xarray as xr from cubed.runtime.create import create_executor +from xarray.namedarray.parallelcompat import list_chunkmanagers +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")) -EXECUTORS = [create_executor("single-threaded"), create_executor("processes")] @pytest.fixture( scope="module", @@ -23,14 +27,14 @@ def executor(request): 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 @@ -40,14 +44,20 @@ def test_chunk(self): 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}) + 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} + 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) \ No newline at end of file + assert_allclose(original, computed) From 576f8fddcb1197473658643a5994b2a62c22e4e1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 7 Aug 2024 09:51:08 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- cubed_xarray/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/cubed_xarray/__init__.py b/cubed_xarray/__init__.py index 01918bd..58b51f5 100644 --- a/cubed_xarray/__init__.py +++ b/cubed_xarray/__init__.py @@ -1,6 +1,5 @@ from importlib.metadata import version - try: __version__ = version("cubed-xarray") except Exception: