diff --git a/changes/3021.feature.rst b/changes/3021.feature.rst new file mode 100644 index 0000000000..8805797ce3 --- /dev/null +++ b/changes/3021.feature.rst @@ -0,0 +1 @@ +Implemented ``move`` for ``LocalStore`` and ``ZipStore``. This allows users to move the store to a different root path. \ No newline at end of file diff --git a/src/zarr/storage/_local.py b/src/zarr/storage/_local.py index f2af75f43e..15b043b1dc 100644 --- a/src/zarr/storage/_local.py +++ b/src/zarr/storage/_local.py @@ -253,5 +253,17 @@ async def list_dir(self, prefix: str) -> AsyncIterator[str]: except (FileNotFoundError, NotADirectoryError): pass + async def move(self, dest_root: Path | str) -> None: + """ + Move the store to another path. The old root directory is deleted. + """ + if isinstance(dest_root, str): + dest_root = Path(dest_root) + os.makedirs(dest_root.parent, exist_ok=True) + if os.path.exists(dest_root): + raise FileExistsError(f"Destination root {dest_root} already exists.") + shutil.move(self.root, dest_root) + self.root = dest_root + async def getsize(self, key: str) -> int: return os.path.getsize(self.root / key) diff --git a/src/zarr/storage/_zip.py b/src/zarr/storage/_zip.py index f9eb8d8808..5d147deded 100644 --- a/src/zarr/storage/_zip.py +++ b/src/zarr/storage/_zip.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import shutil import threading import time import zipfile @@ -288,3 +289,15 @@ async def list_dir(self, prefix: str) -> AsyncIterator[str]: if k not in seen: seen.add(k) yield k + + async def move(self, path: Path | str) -> None: + """ + Move the store to another path. + """ + if isinstance(path, str): + path = Path(path) + self.close() + os.makedirs(path.parent, exist_ok=True) + shutil.move(self.path, path) + self.path = path + await self._open() diff --git a/tests/test_store/test_local.py b/tests/test_store/test_local.py index 8699a85082..7974d0d633 100644 --- a/tests/test_store/test_local.py +++ b/tests/test_store/test_local.py @@ -1,18 +1,18 @@ from __future__ import annotations -from typing import TYPE_CHECKING +import pathlib +import re +import numpy as np import pytest import zarr +from zarr import create_array from zarr.core.buffer import Buffer, cpu from zarr.storage import LocalStore from zarr.testing.store import StoreTests from zarr.testing.utils import assert_bytes_equal -if TYPE_CHECKING: - import pathlib - class TestLocalStore(StoreTests[LocalStore, cpu.Buffer]): store_cls = LocalStore @@ -74,3 +74,38 @@ async def test_get_with_prototype_default(self, store: LocalStore) -> None: await self.set(store, key, data_buf) observed = await store.get(key, prototype=None) assert_bytes_equal(observed, data_buf) + + @pytest.mark.parametrize("ndim", [0, 1, 3]) + @pytest.mark.parametrize( + "destination", ["destination", "foo/bar/destintion", pathlib.Path("foo/bar/destintion")] + ) + async def test_move( + self, tmp_path: pathlib.Path, ndim: int, destination: pathlib.Path | str + ) -> None: + origin = tmp_path / "origin" + if isinstance(destination, str): + destination = str(tmp_path / destination) + else: + destination = tmp_path / destination + + print(type(destination)) + store = await LocalStore.open(root=origin) + shape = (4,) * ndim + chunks = (2,) * ndim + data = np.arange(4**ndim) + if ndim > 0: + data = data.reshape(*shape) + array = create_array(store, data=data, chunks=chunks or "auto") + + await store.move(destination) + + assert store.root == pathlib.Path(destination) + assert pathlib.Path(destination).exists() + assert not origin.exists() + assert np.array_equal(array[...], data) + + store2 = await LocalStore.open(root=origin) + with pytest.raises( + FileExistsError, match=re.escape(f"Destination root {destination} already exists") + ): + await store2.move(destination) diff --git a/tests/test_store/test_zip.py b/tests/test_store/test_zip.py index fa99ca61bd..24b25ed315 100644 --- a/tests/test_store/test_zip.py +++ b/tests/test_store/test_zip.py @@ -10,6 +10,7 @@ import pytest import zarr +from zarr import create_array from zarr.core.buffer import Buffer, cpu, default_buffer_prototype from zarr.core.group import Group from zarr.storage import ZipStore @@ -140,3 +141,17 @@ def test_externally_zipped_store(self, tmp_path: Path) -> None: assert list(zipped.keys()) == list(root.keys()) assert isinstance(group := zipped["foo"], Group) assert list(group.keys()) == list(group.keys()) + + async def test_move(self, tmp_path: Path) -> None: + origin = tmp_path / "origin.zip" + destination = tmp_path / "some_folder" / "destination.zip" + + store = await ZipStore.open(path=origin, mode="a") + array = create_array(store, data=np.arange(10)) + + await store.move(str(destination)) + + assert store.path == destination + assert destination.exists() + assert not origin.exists() + assert np.array_equal(array[...], np.arange(10))