Use async batched _rm() for FsspecStore.delete_dir()

Author: carshadi

src/zarr/storage/_fsspec.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+import inspect
 from typing import TYPE_CHECKING, Any
 
 from zarr.abc.store import ByteRangeRequest, Store
@@ -262,47 +263,22 @@ async def delete(self, key: str) -> None:
             pass
 
     async def delete_dir(self, prefix: str) -> None:
-        """
-        Remove all keys and prefixes in the store that begin with a given prefix.
-        """
+        # docstring inherited
         if not self.supports_deletes:
-            raise NotImplementedError
+            raise NotImplementedError('This method is only available for stores that support deletes.')
         if not self.supports_listing:
-            raise NotImplementedError
+            raise NotImplementedError('This method is only available for stores that support directory listing.')
         self._check_writable()
 
-        if prefix and not prefix.endswith("/"):
-            prefix += "/"
-
-        paths_to_delete = []
-        async for key in self.list_prefix(prefix):
-            paths_to_delete.append(_dereference_path(self.path, key))
+        path_to_delete = _dereference_path(self.path, prefix)
 
-        if not paths_to_delete:
-            return
-
-        try:
-            import s3fs
-        except ImportError:
-            s3fs = None
-
-        # If s3fs is installed and our filesystem is S3FileSystem, do a bulk delete
-        if s3fs and isinstance(self.fs, s3fs.S3FileSystem):
+        if hasattr(self.fs, "_rm") and inspect.iscoroutinefunction(self.fs._rm):
             try:
-                await self.fs._rm(paths_to_delete)
-            except FileNotFoundError:
-                pass
+                await self.fs._rm(path_to_delete, recursive=True)
             except self.allowed_exceptions:
                 pass
         else:
-            # Otherwise, delete one by one
-            for path in paths_to_delete:
-                try:
-                    await self.fs._rm(path)
-                except FileNotFoundError:
-                    pass
-                except self.allowed_exceptions:
-                    pass
+            raise NotImplementedError("The store does not support async deletes")
 
     async def exists(self, key: str) -> bool:
         # docstring inherited