Feature: store learns to delete prefixes when overwriting/creating hierarchy nodes

src/zarr/abc/store.py

@@ -371,6 +371,36 @@ def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         """
         ...
 
+    async def delete_dir(self, prefix: str, recursive: bool = True) -> None:
+        """
+        Remove all keys and prefixes in the store that begin with a given prefix.
+        """
+        if not self.supports_deletes:
+            raise NotImplementedError
+        if not self.supports_listing:
+            raise NotImplementedError
+        self._check_writable()
+        if recursive:
+            if not prefix.endswith("/"):
+                prefix += "/"
+            async for key in self.list_prefix(prefix):
+                await self.delete(key)
+        else:
+            async for key in self.list_dir(prefix):
+                await self.delete(f"{prefix}/{key}")
+
+    async def delete_prefix(self, prefix: str) -> None:
+        """
+        Remove all keys in the store that begin with a given prefix.
+        """
+        if not self.supports_deletes:
+            raise NotImplementedError
+        if not self.supports_listing:
+            raise NotImplementedError
+        self._check_writable()
+        async for key in self.list_prefix(prefix):
+            await self.delete(key)
+
     def close(self) -> None:
         """Close the store."""
         self._is_open = False

src/zarr/core/array.py

@@ -554,7 +554,12 @@ async def _create_v3(
         attributes: dict[str, JSON] | None = None,
         exists_ok: bool = False,
     ) -> AsyncArray[ArrayV3Metadata]:
-        if not exists_ok:
+        if exists_ok:
+            if store_path.store.supports_deletes:
+                await store_path.delete_dir(recursive=True)
+            else:
+                await ensure_no_existing_node(store_path, zarr_format=3)
+        else:
             await ensure_no_existing_node(store_path, zarr_format=3)
 
         shape = parse_shapelike(shape)
@@ -606,7 +611,12 @@ async def _create_v2(
         attributes: dict[str, JSON] | None = None,
         exists_ok: bool = False,
     ) -> AsyncArray[ArrayV2Metadata]:
-        if not exists_ok:
+        if exists_ok:
+            if store_path.store.supports_deletes:
+                await store_path.delete_dir(recursive=True)
+            else:
+                await ensure_no_existing_node(store_path, zarr_format=2)
+        else:
             await ensure_no_existing_node(store_path, zarr_format=2)
 
         if order is None:

src/zarr/core/group.py

@@ -404,7 +404,13 @@ async def from_store(
         zarr_format: ZarrFormat = 3,
     ) -> AsyncGroup:
         store_path = await make_store_path(store)
-        if not exists_ok:
+
+        if exists_ok:
+            if store_path.store.supports_deletes:
+                await store_path.delete_dir(recursive=True)
+            else:
+                await ensure_no_existing_node(store_path, zarr_format=zarr_format)
+        else:
             await ensure_no_existing_node(store_path, zarr_format=zarr_format)
         attributes = attributes or {}
         group = cls(
@@ -710,19 +716,8 @@ def _getitem_consolidated(
 
     async def delitem(self, key: str) -> None:
         store_path = self.store_path / key
-        if self.metadata.zarr_format == 3:
-            await (store_path / ZARR_JSON).delete()
-
-        elif self.metadata.zarr_format == 2:
-            await asyncio.gather(
-                (store_path / ZGROUP_JSON).delete(),  # TODO: missing_ok=False
-                (store_path / ZARRAY_JSON).delete(),  # TODO: missing_ok=False
-                (store_path / ZATTRS_JSON).delete(),  # TODO: missing_ok=True
-            )
-
-        else:
-            raise ValueError(f"unexpected zarr_format: {self.metadata.zarr_format}")
 
+        await store_path.delete_dir(recursive=True)
         if self.metadata.consolidated_metadata:
             self.metadata.consolidated_metadata.metadata.pop(key, None)
             await self._save_metadata()

src/zarr/storage/common.py

@@ -101,6 +101,18 @@ async def delete(self) -> None:
         """
         await self.store.delete(self.path)
 
+    async def delete_dir(self, recursive: bool = False) -> None:
+        """
+        Delete all keys with the given prefix from the store.
+        """
+        await self.store.delete_dir(self.path, recursive=recursive)
+
+    async def delete_prefix(self) -> None:
+        """
+        Delete all keys with the given prefix from the store.
+        """
+        await self.store.delete_prefix(self.path)
+
     async def set_if_not_exists(self, default: Buffer) -> None:
         """
         Store a key to ``value`` if the key is not already present.

src/zarr/storage/local.py

@@ -226,7 +226,8 @@ async def list(self) -> AsyncGenerator[str, None]:
 
     async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
-        to_strip = os.path.join(str(self.root / prefix))
+        to_strip = str(self.root)
+        prefix = prefix.rstrip("/")
         for p in (self.root / prefix).rglob("*"):
             if p.is_file():
                 yield str(p.relative_to(to_strip))

src/zarr/storage/logging.py

@@ -230,3 +230,6 @@ def with_mode(self, mode: AccessModeLiteral) -> Self:
                 log_level=self.log_level,
                 log_handler=self.log_handler,
             )
+
+
+# TODO: wrap delete methods here

src/zarr/storage/memory.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from logging import getLogger
 from typing import TYPE_CHECKING, Self
 
 from zarr.abc.store import ByteRangeRequest, Store
@@ -14,6 +15,9 @@
     from zarr.core.common import AccessModeLiteral
 
 
+logger = getLogger(__name__)
+
+
 class MemoryStore(Store):
     """
     In-memory store for testing purposes.
@@ -137,7 +141,7 @@ async def delete(self, key: str) -> None:
         try:
             del self._store_dict[key]
         except KeyError:
-            pass
+            logger.debug("Key %s does not exist.", key)
 
     async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, bytes]]) -> None:
         # docstring inherited
@@ -150,9 +154,10 @@ async def list(self) -> AsyncGenerator[str, None]:
 
     async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
-        for key in self._store_dict:
+        # note: we materialize all dict keys into a list here so we can mutate the dict in-place (e.g. in delete_prefix)
+        for key in list(self._store_dict):
             if key.startswith(prefix):
-                yield key.removeprefix(prefix)
+                yield key
 
     async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited

src/zarr/storage/remote.py

@@ -340,6 +340,7 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
 
     async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
-        find_str = f"{self.path}/{prefix}"
-        for onefile in await self.fs._find(find_str, detail=False, maxdepth=None, withdirs=False):
-            yield onefile.removeprefix(find_str)
+        for onefile in await self.fs._find(
+            f"{self.path}/{prefix}", detail=False, maxdepth=None, withdirs=False
+        ):
+            yield onefile.removeprefix(f"{self.path}/")

src/zarr/storage/zip.py

@@ -241,7 +241,7 @@ async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
         async for key in self.list():
             if key.startswith(prefix):
-                yield key.removeprefix(prefix)
+                yield key
 
     async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited

src/zarr/testing/store.py

@@ -249,8 +249,7 @@ async def test_list(self, store: S) -> None:
     async def test_list_prefix(self, store: S) -> None:
         """
         Test that the `list_prefix` method works as intended. Given a prefix, it should return
-        all the keys in storage that start with this prefix. Keys should be returned with the shared
-        prefix removed.
+        all the keys in storage that start with this prefix.
         """
         prefixes = ("", "a/", "a/b/", "a/b/c/")
         data = self.buffer_cls.from_bytes(b"")
@@ -264,7 +263,7 @@ async def test_list_prefix(self, store: S) -> None:
             expected: tuple[str, ...] = ()
             for key in store_dict:
                 if key.startswith(prefix):
-                    expected += (key.removeprefix(prefix),)
+                    expected += (key,)
             expected = tuple(sorted(expected))
             assert observed == expected
 

tests/test_array.py

@@ -48,6 +48,9 @@ def test_array_creation_existing_node(
     new_dtype = "float32"
 
     if exists_ok:
+        if not store.supports_deletes:
+            # TODO: confirm you get the expected error here
+            pytest.skip("store does not support deletes")
         arr_new = Array.create(
             spath / "extant",
             shape=new_shape,

tests/test_group.py

@@ -290,6 +290,10 @@ def test_group_open(store: Store, zarr_format: ZarrFormat, exists_ok: bool) -> N
         with pytest.raises(ContainsGroupError):
             Group.from_store(store, attributes=attrs, zarr_format=zarr_format, exists_ok=exists_ok)
     else:
+        if not store.supports_deletes:
+            pytest.skip(
+                "Store does not support deletes but `exists_ok` is True, requiring deletes to override a group"
+            )
         group_created_again = Group.from_store(
             store, attributes=new_attrs, zarr_format=zarr_format, exists_ok=exists_ok
         )
@@ -703,6 +707,8 @@ def test_group_creation_existing_node(
     new_attributes = {"new": True}
 
     if exists_ok:
+        if not store.supports_deletes:
+            pytest.skip("store does not support deletes but exists_ok is True")
         node_new = Group.from_store(
             spath / "extant",
             attributes=new_attributes,
@@ -1075,7 +1081,9 @@ async def test_require_group(store: LocalStore | MemoryStore, zarr_format: ZarrF
     assert foo_group.attrs == {"foo": 100}
 
     # test that we can get the group using require_group and overwrite=True
-    foo_group = await root.require_group("foo", overwrite=True)
+    if store.supports_deletes:
+        foo_group = await root.require_group("foo", overwrite=True)
+        assert foo_group.attrs == {}
 
     _ = await foo_group.create_array(
         "bar", shape=(10,), dtype="uint8", chunk_shape=(2,), attributes={"foo": 100}
@@ -1354,3 +1362,16 @@ def test_group_deprecated_positional_args(method: str) -> None:
     with pytest.warns(FutureWarning, match=r"Pass name=.*, data=.* as keyword args."):
         arr = getattr(root, method)("foo_like", data, **kwargs)
         assert arr.shape == data.shape
+
+
+@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
+def test_delitem_removes_children(store: Store, zarr_format: ZarrFormat) -> None:
+    # https://github.com/zarr-developers/zarr-python/issues/2191
+    g1 = zarr.group(store=store)
+    g1.create_group("0")
+    g1.create_group("0/0")
+    arr = g1.create_array("0/0/0", shape=(1,))
+    arr[:] = 1
+    del g1["0"]
+    with pytest.raises(KeyError):
+        g1["0/0"]

tests/test_store/test_logging.py

@@ -45,10 +45,14 @@ async def test_logging_store_counter(store: Store) -> None:
     arr[:] = 1
 
     assert wrapped.counter["set"] == 2
-    assert wrapped.counter["get"] == 0  # 1 if overwrite=False
     assert wrapped.counter["list"] == 0
     assert wrapped.counter["list_dir"] == 0
-    assert wrapped.counter["list_prefix"] == 0
+    if store.supports_deletes:
+        assert wrapped.counter["get"] == 0  # 1 if overwrite=False
+        assert wrapped.counter["list_prefix"] == 1
+    else:
+        assert wrapped.counter["get"] == 1
+        assert wrapped.counter["list_prefix"] == 0
 
 
 async def test_with_mode():

tests/test_store/test_zip.py

@@ -103,4 +103,4 @@ async def test_with_mode(self, store: ZipStore) -> None:
 
     @pytest.mark.parametrize("mode", ["a", "w"])
     async def test_store_open_mode(self, store_kwargs: dict[str, Any], mode: str) -> None:
-        super().test_store_open_mode(store_kwargs, mode)
+        await super().test_store_open_mode(store_kwargs, mode)