From c38c2f7cd52dfa884e4f542de6aaed73a57835c7 Mon Sep 17 00:00:00 2001
From: Joseph Hamman <joe@earthmover.io>
Date: Thu, 24 Oct 2024 16:29:35 -0700
Subject: [PATCH] respond to review

---
 src/zarr/abc/store.py            | 26 +++++---------------------
 src/zarr/core/array.py           |  4 ++--
 src/zarr/core/group.py           |  4 ++--
 src/zarr/storage/common.py       | 13 +++++--------
 src/zarr/storage/logging.py      |  5 +++++
 src/zarr/testing/store.py        | 13 +++++++++++++
 tests/test_store/test_logging.py |  5 +++--
 tests/test_store/test_zip.py     |  4 ----
 8 files changed, 35 insertions(+), 39 deletions(-)

diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py
index 03f42584f7..c7f21df508 100644
--- a/src/zarr/abc/store.py
+++ b/src/zarr/abc/store.py
@@ -342,8 +342,8 @@ def list(self) -> AsyncGenerator[str, None]:
     @abstractmethod
     def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         """
-        Retrieve all keys in the store that begin with a given prefix. Keys are returned with the
-        common leading prefix removed.
+        Retrieve all keys in the store that begin with a given prefix. Keys are returned as
+        absolute paths (i.e. including the prefix).
 
         Parameters
         ----------
@@ -371,7 +371,7 @@ def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         """
         ...
 
-    async def delete_dir(self, prefix: str, recursive: bool = True) -> None:
+    async def delete_dir(self, prefix: str) -> None:
         """
         Remove all keys and prefixes in the store that begin with a given prefix.
         """
@@ -380,24 +380,8 @@ async def delete_dir(self, prefix: str, recursive: bool = True) -> None:
         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()
+        if not prefix.endswith("/"):
+            prefix += "/"
         async for key in self.list_prefix(prefix):
             await self.delete(key)
 
diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py
index 4549cc3e7a..7b5733116c 100644
--- a/src/zarr/core/array.py
+++ b/src/zarr/core/array.py
@@ -556,7 +556,7 @@ async def _create_v3(
     ) -> AsyncArray[ArrayV3Metadata]:
         if exists_ok:
             if store_path.store.supports_deletes:
-                await store_path.delete_dir(recursive=True)
+                await store_path.delete_dir()
             else:
                 await ensure_no_existing_node(store_path, zarr_format=3)
         else:
@@ -613,7 +613,7 @@ async def _create_v2(
     ) -> AsyncArray[ArrayV2Metadata]:
         if exists_ok:
             if store_path.store.supports_deletes:
-                await store_path.delete_dir(recursive=True)
+                await store_path.delete_dir()
             else:
                 await ensure_no_existing_node(store_path, zarr_format=2)
         else:
diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py
index 914f001cce..5e68039305 100644
--- a/src/zarr/core/group.py
+++ b/src/zarr/core/group.py
@@ -407,7 +407,7 @@ async def from_store(
 
         if exists_ok:
             if store_path.store.supports_deletes:
-                await store_path.delete_dir(recursive=True)
+                await store_path.delete_dir()
             else:
                 await ensure_no_existing_node(store_path, zarr_format=zarr_format)
         else:
@@ -717,7 +717,7 @@ def _getitem_consolidated(
     async def delitem(self, key: str) -> None:
         store_path = self.store_path / key
 
-        await store_path.delete_dir(recursive=True)
+        await store_path.delete_dir()
         if self.metadata.consolidated_metadata:
             self.metadata.consolidated_metadata.metadata.pop(key, None)
             await self._save_metadata()
diff --git a/src/zarr/storage/common.py b/src/zarr/storage/common.py
index af434ecd34..337fbc59a4 100644
--- a/src/zarr/storage/common.py
+++ b/src/zarr/storage/common.py
@@ -101,17 +101,14 @@ async def delete(self) -> None:
         """
         await self.store.delete(self.path)
 
-    async def delete_dir(self, recursive: bool = False) -> None:
+    async def delete_dir(self) -> 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)
+        path = self.path
+        if not path.endswith("/"):
+            path += "/"
+        await self.store.delete_dir(path)
 
     async def set_if_not_exists(self, default: Buffer) -> None:
         """
diff --git a/src/zarr/storage/logging.py b/src/zarr/storage/logging.py
index 453a793cbc..755d7a58e6 100644
--- a/src/zarr/storage/logging.py
+++ b/src/zarr/storage/logging.py
@@ -222,6 +222,11 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
             async for key in self._store.list_dir(prefix=prefix):
                 yield key
 
+    async def delete_dir(self, prefix: str) -> None:
+        # docstring inherited
+        with self.log(prefix):
+            await self._store.delete_dir(prefix=prefix)
+
     def with_mode(self, mode: AccessModeLiteral) -> Self:
         # docstring inherited
         with self.log(mode):
diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py
index fd9d6e6910..b39c548fbe 100644
--- a/src/zarr/testing/store.py
+++ b/src/zarr/testing/store.py
@@ -213,11 +213,24 @@ async def test_exists(self, store: S) -> None:
         assert await store.exists("foo/zarr.json")
 
     async def test_delete(self, store: S) -> None:
+        if not store.supports_deletes:
+            pytest.skip("store does not support deletes")
         await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar"))
         assert await store.exists("foo/zarr.json")
         await store.delete("foo/zarr.json")
         assert not await store.exists("foo/zarr.json")
 
+    async def test_delete_dir(self, store: S) -> None:
+        if not store.supports_deletes:
+            pytest.skip("store does not support deletes")
+        await store.set("zarr.json", self.buffer_cls.from_bytes(b"root"))
+        await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar"))
+        await store.set("foo/c/0", self.buffer_cls.from_bytes(b"chunk"))
+        await store.delete_dir("foo")
+        assert await store.exists("zarr.json")
+        assert not await store.exists("foo/zarr.json")
+        assert not await store.exists("foo/c/0")
+
     async def test_empty(self, store: S) -> None:
         assert await store.empty()
         await self.set(
diff --git a/tests/test_store/test_logging.py b/tests/test_store/test_logging.py
index 976da68185..50db5c1c5b 100644
--- a/tests/test_store/test_logging.py
+++ b/tests/test_store/test_logging.py
@@ -47,12 +47,13 @@ async def test_logging_store_counter(store: Store) -> None:
     assert wrapped.counter["set"] == 2
     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
+        assert wrapped.counter["delete_dir"] == 1
     else:
         assert wrapped.counter["get"] == 1
-        assert wrapped.counter["list_prefix"] == 0
+        assert wrapped.counter["delete_dir"] == 0
 
 
 async def test_with_mode():
diff --git a/tests/test_store/test_zip.py b/tests/test_store/test_zip.py
index eb129a1b2d..8dee474498 100644
--- a/tests/test_store/test_zip.py
+++ b/tests/test_store/test_zip.py
@@ -14,7 +14,6 @@
 from zarr.testing.store import StoreTests
 
 if TYPE_CHECKING:
-    from collections.abc import Coroutine
     from typing import Any
 
 
@@ -65,9 +64,6 @@ def test_store_supports_partial_writes(self, store: ZipStore) -> None:
     def test_store_supports_listing(self, store: ZipStore) -> None:
         assert store.supports_listing
 
-    def test_delete(self, store: ZipStore) -> Coroutine[Any, Any, None]:
-        pass
-
     def test_api_integration(self, store: ZipStore) -> None:
         root = zarr.open_group(store=store)