obstore implementations for .getsize and .getsize_prefix

Author: slowjazz

src/zarr/storage/_obstore.py

@@ -4,7 +4,7 @@
 import contextlib
 import pickle
 from collections import defaultdict
-from typing import TYPE_CHECKING, TypedDict
+from typing import TYPE_CHECKING, TypedDict, TypeVar
 
 from zarr.abc.store import (
     ByteRequest,
@@ -16,7 +16,7 @@
 from zarr.core.config import config
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Coroutine, Iterable
+    from collections.abc import AsyncGenerator, Coroutine, Iterable, Sequence
     from typing import Any
 
     from obstore import ListResult, ListStream, ObjectMeta, OffsetRange, SuffixRange
@@ -27,6 +27,8 @@
 
 __all__ = ["ObjectStore"]
 
+T = TypeVar("T")
+
 _ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = (
     FileNotFoundError,
     IsADirectoryError,
@@ -212,41 +214,55 @@ def supports_listing(self) -> bool:
         # docstring inherited
         return True
 
-    def list(self) -> AsyncGenerator[str, None]:
-        # docstring inherited
+    def _list(self, prefix: str | None = None) -> AsyncGenerator[ObjectMeta, None]:
         import obstore as obs
 
-        objects: ListStream[list[ObjectMeta]] = obs.list(self.store)
-        return _transform_list(objects)
+        objects: ListStream[Sequence[ObjectMeta]] = obs.list(self.store, prefix=prefix)
+        return (obj async for obj in _transform_list(objects))
 
-    def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
+    def list(self) -> AsyncGenerator[str, None]:
         # docstring inherited
-        import obstore as obs
+        return (obj["path"] async for obj in self._list())
 
-        objects: ListStream[list[ObjectMeta]] = obs.list(self.store, prefix=prefix)
-        return _transform_list(objects)
+    def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
+        # docstring inherited
+        return (obj["path"] async for obj in self._list(prefix))
 
     def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
         import obstore as obs
 
-        coroutine = obs.list_with_delimiter_async(self.store, prefix=prefix)
+        coroutine: Coroutine[Any, Any, ListResult[Sequence[ObjectMeta]]] = (
+            obs.list_with_delimiter_async(self.store, prefix=prefix)
+        )
         return _transform_list_dir(coroutine, prefix)
 
+    async def getsize(self, key: str) -> int:
+        # docstring inherited
+        import obstore as obs
+
+        resp = await obs.head_async(self.store, key)
+        return resp["size"]
+
+    async def getsize_prefix(self, prefix: str) -> int:
+        # docstring inherited
+        sizes = [obj["size"] async for obj in self._list(prefix=prefix)]
+        return sum(sizes)
+
 
 async def _transform_list(
-    list_stream: ListStream[list[ObjectMeta]],
-) -> AsyncGenerator[str, None]:
+    list_stream: ListStream[Sequence[ObjectMeta]],
+) -> AsyncGenerator[ObjectMeta, None]:
     """
-    Transform the result of list into an async generator of paths.
+    Transform the result of list into an async generator of ObjectMeta dicts.
     """
     async for batch in list_stream:
         for item in batch:
-            yield item["path"]
+            yield item
 
 
 async def _transform_list_dir(
-    list_result_coroutine: Coroutine[Any, Any, ListResult[list[ObjectMeta]]], prefix: str
+    list_result_coroutine: Coroutine[Any, Any, ListResult[Sequence[ObjectMeta]]], prefix: str
 ) -> AsyncGenerator[str, None]:
     """
     Transform the result of list_with_delimiter into an async generator of paths.

tests/test_store/test_object.py

@@ -75,6 +75,21 @@ def test_store_init_raises(self) -> None:
         with pytest.raises(TypeError):
             ObjectStore("path/to/store")
 
+    async def test_store_getsize(self, store: ObjectStore) -> None:
+        buf = cpu.Buffer.from_bytes(b"\x01\x02\x03\x04")
+        await self.set(store, "key", buf)
+        size = await store.getsize("key")
+        assert size == len(buf)
+
+    async def test_store_getsize_prefix(self, store: ObjectStore) -> None:
+        buf = cpu.Buffer.from_bytes(b"\x01\x02\x03\x04")
+        await self.set(store, "c/key1/0", buf)
+        await self.set(store, "c/key2/0", buf)
+        size = await store.getsize_prefix("c/key1")
+        assert size == len(buf)
+        total_size = await store.getsize_prefix("c")
+        assert total_size == len(buf) * 2
+
 
 @pytest.mark.slow_hypothesis
 def test_zarr_hierarchy():