Merge pull request #6 from kylebarron/kyle/fix-transform-list-dir

Use `list_with_delimiter` in `list_dir` method

Author: kylebarron

pyproject.toml

@@ -86,7 +86,6 @@ test = [
     "mypy",
     "hypothesis",
     "universal-pathlib",
-    "obstore==0.3.0",
 ]
 optional = ["rich", "universal-pathlib"]
 docs = [
@@ -213,6 +212,7 @@ dependencies = [
     'universal_pathlib @ git+https://github.com/fsspec/universal_pathlib',
     'typing_extensions @ git+https://github.com/python/typing_extensions',
     'donfig @ git+https://github.com/pytroll/donfig',
+    "obstore @ git+https://github.com/developmentseed/obstore@main#subdirectory=obstore",
     # test deps
     'hypothesis',
     'pytest',

src/zarr/storage/object_store.py

@@ -22,7 +22,7 @@
     from collections.abc import AsyncGenerator, Coroutine, Iterable
     from typing import Any
 
-    from obstore import ListStream, ObjectMeta, OffsetRange, SuffixRange
+    from obstore import ListResult, ListStream, ObjectMeta, OffsetRange, SuffixRange
     from obstore.store import ObjectStore as _ObjectStore
 
     from zarr.core.buffer import Buffer, BufferPrototype
@@ -168,8 +168,8 @@ def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         return _transform_list(objects)
 
     def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
-        objects: ListStream[list[ObjectMeta]] = obs.list(self.store, prefix=prefix)
-        return _transform_list_dir(objects, prefix)
+        coroutine = obs.list_with_delimiter_async(self.store, prefix=prefix)
+        return _transform_list_dir(coroutine, prefix)
 
 
 async def _transform_list(
@@ -180,19 +180,23 @@ async def _transform_list(
             yield item["path"]
 
 
+# ['zarr.json', 'c/0', 'c/1'] -> ['zarr.json']
+# ['zarr.json', 'c/0', 'c/1'] -> ['zarr.json, 'c']
 async def _transform_list_dir(
-    list_stream: AsyncGenerator[list[ObjectMeta], None], prefix: str
+    list_result_coroutine: Coroutine[Any, Any, ListResult], prefix: str
 ) -> AsyncGenerator[str, None]:
+    """
+    Transform the result of list_with_delimiter into an async generator of prefixes and paths.
+    """
+    list_result = await list_result_coroutine
+
     # We assume that the underlying object-store implementation correctly handles the
     # prefix, so we don't double-check that the returned results actually start with the
     # given prefix.
-    prefix_len = len(prefix) + 1  # If one is not added to the length, all items will contain "/"
-    async for batch in list_stream:
-        for item in batch:
-            # Yield this item if "/" does not exist after the prefix
-            item_path = item["path"][prefix_len:]
-            if "/" not in item_path:
-                yield item_path
+    prefixes = list_result["common_prefixes"]
+    objects = [obj["path"].lstrip(prefix) for obj in list_result["objects"]]
+    for item in prefixes + objects:
+        yield item
 
 
 class _BoundedRequest(TypedDict):