Feat: improves delete_dir for s3fs-backed FsspecStore

src/zarr/storage/_fsspec.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import warnings
+import inspect
+from contextlib import suppress
 from typing import TYPE_CHECKING, Any
 
 from zarr.abc.store import (
@@ -281,6 +283,19 @@
         except self.allowed_exceptions:
             pass
 
+    async def delete_dir(self, prefix: str) -> None:
+        # docstring inherited
+        if not self.supports_deletes:
+            raise NotImplementedError('This method is only available for stores that support deletes.')
+        if not self.supports_listing:
+            raise NotImplementedError('This method is only available for stores that support directory listing.')
+        self._check_writable()
+
+        path_to_delete = _dereference_path(self.path, prefix)
+
+        with suppress(*self.allowed_exceptions):
+            await self.fs._rm(path_to_delete, recursive=True)
+
     async def exists(self, key: str) -> bool:
         # docstring inherited
         path = _dereference_path(self.path, key)

tests/test_store/test_fsspec.py

@@ -217,6 +217,20 @@ async def test_empty_nonexistent_path(self, store_kwargs) -> None:
         store = await self.store_cls.open(**store_kwargs)
         assert await store.is_empty("")
 
+    async def test_delete_dir_unsupported_deletes(self, store: FsspecStore) -> None:
+        store.supports_deletes = False
+        with pytest.raises(
+            NotImplementedError, match=".*only available for stores that support deletes."
+        ):
+            await store.delete_dir("test_prefix")
+
+    async def test_delete_dir_unsupported_listing(self, store: FsspecStore) -> None:
+        store.supports_listing = False
+        with pytest.raises(
+            NotImplementedError, match=".*only available for stores that support directory listing."
+        ):
+            await store.delete_dir("test_prefix")
+
 
 @pytest.mark.skipif(
     parse_version(fsspec.__version__) < parse_version("2024.12.0"),
@@ -244,3 +258,27 @@ def test_no_wrap_async_filesystem():
 
     assert not isinstance(store.fs, AsyncFileSystemWrapper)
     assert store.fs.async_impl
+
+
+@pytest.mark.skipif(
+    parse_version(fsspec.__version__) < parse_version("2024.12.0"),
+    reason="No AsyncFileSystemWrapper",
+)
+async def test_delete_dir_wrapped_filesystem(tmpdir) -> None:
+    """The local fs is not async so we should expect it to be wrapped automatically"""
+    from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
+
+    store = FsspecStore.from_url(tmpdir / "test/path", storage_options={"auto_mkdir": True})
+
+    assert isinstance(store.fs, AsyncFileSystemWrapper)
+    assert store.fs.async_impl
+
+    await store.set("zarr.json", cpu.Buffer.from_bytes(b"root"))
+    await store.set("foo-bar/zarr.json", cpu.Buffer.from_bytes(b"root"))
+    await store.set("foo/zarr.json", cpu.Buffer.from_bytes(b"bar"))
+    await store.set("foo/c/0", cpu.Buffer.from_bytes(b"chunk"))
+    await store.delete_dir("foo")
+    assert await store.exists("zarr.json")
+    assert await store.exists("foo-bar/zarr.json")
+    assert not await store.exists("foo/zarr.json")
+    assert not await store.exists("foo/c/0")