Use prefix

Author: TomAugspurger

src/zarr/abc/store.py

@@ -417,9 +417,9 @@ async def getsize(self, key: str) -> int:
             raise FileNotFoundError(key)
         return len(value)
 
-    async def getsize_dir(self, prefix: str) -> int:
+    async def getsize_prefix(self, prefix: str) -> int:
         """
-        Return the size, in bytes, of all values in a directory.
+        Return the size, in bytes, of all values under a prefix.
 
         This will include just values whose keys start with ``prefix`` and
         do not contain the character ``/`` after the given prefix.
@@ -440,7 +440,7 @@ async def getsize_dir(self, prefix: str) -> int:
         listing all the keys in a directory and calling :meth:`Store.getsize`
         on each.
         """
-        keys = [x async for x in self.list_dir(prefix)]
+        keys = [x async for x in self.list_prefix(prefix)]
         sizes = await gather(*[self.getsize(key) for key in keys])
         return sum(sizes)
 

src/zarr/core/array.py

@@ -844,7 +844,7 @@ def nchunks_initialized(self) -> int:
         return nchunks_initialized(self)
 
     async def nbytes_stored(self) -> int:
-        return await self.store_path.store.getsize_dir(self.store_path.path)
+        return await self.store_path.store.getsize_prefix(self.store_path.path)
 
     def _iter_chunk_coords(
         self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None

tests/test_array.py

@@ -371,13 +371,28 @@ def test_chunks_initialized(test_cls: type[Array] | type[AsyncArray[Any]]) -> No
         assert observed == expected
 
 
-def test_nbytes_stored():
+def test_nbytes_stored() -> None:
     arr = zarr.create(shape=(100,), chunks=(10,), dtype="i4")
     result = arr.nbytes_stored()
-    assert result == 366  # the size of the metadata document. This is a fragile test
+    assert result == 366  # the size of the metadata document. This is a fragile test.
     arr[:50] = 1
     result = arr.nbytes_stored()
-    assert result == 366
+    assert result == 566  # the size with 5 chunks filled.
+    arr[50:] = 2
+    result = arr.nbytes_stored()
+    assert result == 766  # the size with all chunks filled.
+
+
+async def test_nbytes_stored_async() -> None:
+    arr = await zarr.api.asynchronous.create(shape=(100,), chunks=(10,), dtype="i4")
+    result = await arr.nbytes_stored()
+    assert result == 366  # the size of the metadata document. This is a fragile test.
+    await arr.setitem(slice(50), 1)
+    result = await arr.nbytes_stored()
+    assert result == 566  # the size with 5 chunks filled.
+    await arr.setitem(slice(50, 100), 2)
+    result = await arr.nbytes_stored()
+    assert result == 766  # the size with all chunks filled.
 
 
 def test_default_fill_values() -> None: