Merge branch 'main' of https://github.com/zarr-developers/zarr-python into refactor/store-mode

Author: jhamman

src/zarr/abc/store.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Iterable
+    from collections.abc import AsyncGenerator, AsyncIterator, Iterable
     from types import TracebackType
     from typing import Any, Self, TypeAlias
 
@@ -263,16 +263,19 @@ def supports_listing(self) -> bool:
         ...
 
     @abstractmethod
-    def list(self) -> AsyncGenerator[str]:
+    def list(self) -> AsyncIterator[str]:
         """Retrieve all keys in the store.
 
         Returns
         -------
-        AsyncGenerator[str, None]
+        AsyncIterator[str]
         """
+        # This method should be async, like overridden methods in child classes.
+        # However, that's not straightforward:
+        # https://stackoverflow.com/questions/68905848
 
     @abstractmethod
-    def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         """
         Retrieve all keys in the store that begin with a given prefix. Keys are returned relative
         to the root of the store.
@@ -283,11 +286,14 @@ def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
 
         Returns
         -------
-        AsyncGenerator[str, None]
+        AsyncIterator[str]
         """
+        # This method should be async, like overridden methods in child classes.
+        # However, that's not straightforward:
+        # https://stackoverflow.com/questions/68905848
 
     @abstractmethod
-    def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    def list_dir(self, prefix: str) -> AsyncIterator[str]:
         """
         Retrieve all keys and prefixes with a given prefix and which do not contain the character
         “/” after the given prefix.
@@ -298,8 +304,11 @@ def list_dir(self, prefix: str) -> AsyncGenerator[str]:
 
         Returns
         -------
-        AsyncGenerator[str, None]
+        AsyncIterator[str]
         """
+        # This method should be async, like overridden methods in child classes.
+        # However, that's not straightforward:
+        # https://stackoverflow.com/questions/68905848
 
     async def delete_dir(self, prefix: str) -> None:
         """

src/zarr/core/array.py

@@ -645,6 +645,29 @@ def from_dict(
         store_path: StorePath,
         data: dict[str, JSON],
     ) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]:
+        """
+        Create a Zarr array from a dictionary, with support for both Zarr v2 and v3 metadata.
+
+        Parameters
+        ----------
+        store_path : StorePath
+            The path within the store where the array should be created.
+
+        data : dict
+            A dictionary representing the array data. This dictionary should include necessary metadata
+            for the array, such as shape, dtype, and other attributes. The format of the metadata
+            will determine whether a Zarr v2 or v3 array is created.
+
+        Returns
+        -------
+        AsyncArray[ArrayV3Metadata] or AsyncArray[ArrayV2Metadata]
+            The created Zarr array, either using v2 or v3 metadata based on the provided data.
+
+        Raises
+        ------
+        ValueError
+            If the dictionary data is invalid or incompatible with either Zarr v2 or v3 array creation.
+        """
         metadata = parse_array_metadata(data)
         return cls(metadata=metadata, store_path=store_path)
 
@@ -1040,6 +1063,9 @@ async def getitem(
         return await self._get_selection(indexer, prototype=prototype)
 
     async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = False) -> None:
+        """
+        Asynchronously save the array metadata.
+        """
         to_save = metadata.to_buffer_dict(default_buffer_prototype())
         awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()]
 
@@ -1118,6 +1144,39 @@ async def setitem(
         value: npt.ArrayLike,
         prototype: BufferPrototype | None = None,
     ) -> None:
+        """
+        Asynchronously set values in the array using basic indexing.
+
+        Parameters
+        ----------
+        selection : BasicSelection
+            The selection defining the region of the array to set.
+
+        value : numpy.typing.ArrayLike
+            The values to be written into the selected region of the array.
+
+        prototype : BufferPrototype or None, optional
+            A prototype buffer that defines the structure and properties of the array chunks being modified.
+            If None, the default buffer prototype is used. Default is None.
+
+        Returns
+        -------
+        None
+            This method does not return any value.
+
+        Raises
+        ------
+        IndexError
+            If the selection is out of bounds for the array.
+
+        ValueError
+            If the values are not compatible with the array's dtype or shape.
+
+        Notes
+        -----
+        - This method is asynchronous and should be awaited.
+        - Supports basic indexing, where the selection is contiguous and does not involve advanced indexing.
+        """
         if prototype is None:
             prototype = default_buffer_prototype()
         indexer = BasicIndexer(
@@ -1128,6 +1187,32 @@ async def setitem(
         return await self._set_selection(indexer, value, prototype=prototype)
 
     async def resize(self, new_shape: ShapeLike, delete_outside_chunks: bool = True) -> None:
+        """
+        Asynchronously resize the array to a new shape.
+
+        Parameters
+        ----------
+        new_shape : ChunkCoords
+            The desired new shape of the array.
+
+        delete_outside_chunks : bool, optional
+            If True (default), chunks that fall outside the new shape will be deleted. If False,
+            the data in those chunks will be preserved.
+
+        Returns
+        -------
+        AsyncArray
+            The resized array.
+
+        Raises
+        ------
+        ValueError
+            If the new shape is incompatible with the current array's chunking configuration.
+
+        Notes
+        -----
+        - This method is asynchronous and should be awaited.
+        """
         new_shape = parse_shapelike(new_shape)
         assert len(new_shape) == len(self.metadata.shape)
         new_metadata = self.metadata.update_shape(new_shape)
@@ -1210,6 +1295,31 @@ async def append(self, data: npt.ArrayLike, axis: int = 0) -> ChunkCoords:
         return new_shape
 
     async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self:
+        """
+        Asynchronously update the array's attributes.
+
+        Parameters
+        ----------
+        new_attributes : dict of str to JSON
+            A dictionary of new attributes to update or add to the array. The keys represent attribute
+            names, and the values must be JSON-compatible.
+
+        Returns
+        -------
+        AsyncArray
+            The array with the updated attributes.
+
+        Raises
+        ------
+        ValueError
+            If the attributes are invalid or incompatible with the array's metadata.
+
+        Notes
+        -----
+        - This method is asynchronous and should be awaited.
+        - The updated attributes will be merged with existing attributes, and any conflicts will be
+          overwritten by the new values.
+        """
         # metadata.attributes is "frozen" so we simply clear and update the dict
         self.metadata.attributes.clear()
         self.metadata.attributes.update(new_attributes)
@@ -1328,6 +1438,28 @@ def from_dict(
         store_path: StorePath,
         data: dict[str, JSON],
     ) -> Array:
+        """
+        Create a Zarr array from a dictionary.
+
+        Parameters
+        ----------
+        store_path : StorePath
+            The path within the store where the array should be created.
+
+        data : dict
+            A dictionary representing the array data. This dictionary should include necessary metadata
+            for the array, such as shape, dtype, fill value, and attributes.
+
+        Returns
+        -------
+        Array
+            The created Zarr array.
+
+        Raises
+        ------
+        ValueError
+            If the dictionary data is invalid or missing required fields for array creation.
+        """
         async_array = AsyncArray.from_dict(store_path=store_path, data=data)
         return cls(async_array)
 
@@ -2934,6 +3066,30 @@ def append(self, data: npt.ArrayLike, axis: int = 0) -> ChunkCoords:
         return sync(self._async_array.append(data, axis=axis))
 
     def update_attributes(self, new_attributes: dict[str, JSON]) -> Array:
+        """
+        Update the array's attributes.
+
+        Parameters
+        ----------
+        new_attributes : dict
+            A dictionary of new attributes to update or add to the array. The keys represent attribute
+            names, and the values must be JSON-compatible.
+
+        Returns
+        -------
+        Array
+            The array with the updated attributes.
+
+        Raises
+        ------
+        ValueError
+            If the attributes are invalid or incompatible with the array's metadata.
+
+        Notes
+        -----
+        - The updated attributes will be merged with existing attributes, and any conflicts will be
+          overwritten by the new values.
+        """
         # TODO: remove this cast when type inference improves
         new_array = sync(self._async_array.update_attributes(new_attributes))
         # TODO: remove this cast when type inference improves

src/zarr/storage/local.py

@@ -11,7 +11,7 @@
 from zarr.core.common import concurrent_map
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Iterable
+    from collections.abc import AsyncIterator, Iterable
 
     from zarr.core.buffer import BufferPrototype
 
@@ -198,22 +198,22 @@ async def exists(self, key: str) -> bool:
         path = self.root / key
         return await asyncio.to_thread(path.is_file)
 
-    async def list(self) -> AsyncGenerator[str]:
+    async def list(self) -> AsyncIterator[str]:
         # docstring inherited
         to_strip = self.root.as_posix() + "/"
         for p in list(self.root.rglob("*")):
             if p.is_file():
                 yield p.as_posix().replace(to_strip, "")
 
-    async def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         to_strip = self.root.as_posix() + "/"
         prefix = prefix.rstrip("/")
         for p in (self.root / prefix).rglob("*"):
             if p.is_file():
                 yield p.as_posix().replace(to_strip, "")
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         base = self.root / prefix
         try:

src/zarr/storage/logging.py

@@ -10,7 +10,7 @@
 from zarr.abc.store import ByteRangeRequest, Store
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Generator, Iterable
+    from collections.abc import AsyncIterator, Generator, Iterable
 
     from zarr.core.buffer import Buffer, BufferPrototype
 
@@ -203,19 +203,19 @@ async def set_partial_values(
         with self.log(keys):
             return await self._store.set_partial_values(key_start_values=key_start_values)
 
-    async def list(self) -> AsyncGenerator[str]:
+    async def list(self) -> AsyncIterator[str]:
         # docstring inherited
         with self.log():
             async for key in self._store.list():
                 yield key
 
-    async def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         with self.log(prefix):
             async for key in self._store.list_prefix(prefix=prefix):
                 yield key
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         with self.log(prefix):
             async for key in self._store.list_dir(prefix=prefix):

src/zarr/storage/memory.py

@@ -9,7 +9,7 @@
 from zarr.storage._utils import _normalize_interval_index
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Iterable, MutableMapping
+    from collections.abc import AsyncIterator, Iterable, MutableMapping
 
     from zarr.core.buffer import BufferPrototype
 
@@ -138,19 +138,19 @@ async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, by
         # docstring inherited
         raise NotImplementedError
 
-    async def list(self) -> AsyncGenerator[str]:
+    async def list(self) -> AsyncIterator[str]:
         # docstring inherited
         for key in self._store_dict:
             yield key
 
-    async def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         # note: we materialize all dict keys into a list here so we can mutate the dict in-place (e.g. in delete_prefix)
         for key in list(self._store_dict):
             if key.startswith(prefix):
                 yield key
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         prefix = prefix.rstrip("/")
 

src/zarr/storage/remote.py

@@ -7,7 +7,7 @@
 from zarr.storage.common import _dereference_path
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Iterable
+    from collections.abc import AsyncIterator, Iterable
 
     from fsspec.asyn import AsyncFileSystem
 
@@ -303,13 +303,13 @@ async def set_partial_values(
         # docstring inherited
         raise NotImplementedError
 
-    async def list(self) -> AsyncGenerator[str]:
+    async def list(self) -> AsyncIterator[str]:
         # docstring inherited
         allfiles = await self.fs._find(self.path, detail=False, withdirs=False)
         for onefile in (a.replace(self.path + "/", "") for a in allfiles):
             yield onefile
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         prefix = f"{self.path}/{prefix.rstrip('/')}"
         try:
@@ -319,7 +319,7 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
         for onefile in (a.replace(prefix + "/", "") for a in allfiles):
             yield onefile.removeprefix(self.path).removeprefix("/")
 
-    async def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         for onefile in await self.fs._find(
             f"{self.path}/{prefix}", detail=False, maxdepth=None, withdirs=False