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 85d244f17b..bf7c20702b 100644 --- a/src/zarr/storage/_local.py +++ b/src/zarr/storage/_local.py @@ -253,5 +253,21 @@ 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, exist_ok=True) + for src_file in self.root.rglob("*"): + if src_file.is_file(): + relative_path = src_file.relative_to(self.root) + dest_file_path = dest_root / relative_path + dest_file_path.parent.mkdir(parents=True, exist_ok=True) + shutil.move(str(src_file), str(dest_file_path)) + shutil.rmtree(self.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 d9d941c6f0..0af990aecc 100644 --- a/tests/test_store/test_local.py +++ b/tests/test_store/test_local.py @@ -2,9 +2,11 @@ from typing import TYPE_CHECKING +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 @@ -74,3 +76,23 @@ async def test_get_with_prototype_default(self, store: LocalStore): 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, 2, 3]) + async def test_move(self, tmp_path: pathlib.Path, ndim): + origin = tmp_path / "origin" + destination = tmp_path / "destintion" + + 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(str(destination)) + + assert store.root == destination + assert destination.exists() + assert not origin.exists() + assert np.array_equal(array[...], data) diff --git a/tests/test_store/test_zip.py b/tests/test_store/test_zip.py index 0237258ab1..760fbe4a25 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.storage import ZipStore from zarr.testing.store import StoreTests @@ -135,3 +136,17 @@ def test_externally_zipped_store(self, tmp_path: Path) -> None: zipped = zarr.open_group(ZipStore(zip_path, mode="r"), mode="r") assert list(zipped.keys()) == list(root.keys()) assert list(zipped["foo"].keys()) == list(root["foo"].keys()) + + async def test_move(self, tmp_path: Path): + 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))