Implement Store.move

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.

src/zarr/storage/_local.py

@@ -253,5 +253,21 @@
         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)

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 @@
                     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()

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)

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))