Implement Store.move (#3021)

* move for LocalStore

* move for ZipStore

* remove redundant check

* open and close zipstore

* fix zipstore.move

* fix localstore.move for ndim>1

* format

* remove abstract Store .move

* document changes

* shutil.move entire folder

* fix test for windows

Author: brokkoli71

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

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

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)

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