clean up array config parsing, and modify init_array to create parent groups and return an asyncarray instead of metadata

Author: d-v-b

src/zarr/api/asynchronous.py

@@ -10,7 +10,7 @@
 from typing_extensions import deprecated
 
 from zarr.core.array import Array, AsyncArray, create_array, get_array_metadata
-from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
+from zarr.core.array_spec import ArrayConfig, ArrayConfigLike, ArrayConfigParams
 from zarr.core.buffer import NDArrayLike
 from zarr.core.common import (
     JSON,
@@ -856,7 +856,7 @@ async def create(
     codecs: Iterable[Codec | dict[str, JSON]] | None = None,
     dimension_names: Iterable[str] | None = None,
     storage_options: dict[str, Any] | None = None,
-    config: ArrayConfig | ArrayConfigLike | None = None,
+    config: ArrayConfigLike | None = None,
     **kwargs: Any,
 ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array.
@@ -1018,7 +1018,7 @@ async def create(
         mode = "a"
     store_path = await make_store_path(store, path=path, mode=mode, storage_options=storage_options)
 
-    config_dict: ArrayConfigLike = {}
+    config_dict: ArrayConfigParams = {}
 
     if write_empty_chunks is not None:
         if config is not None:

src/zarr/api/synchronous.py

@@ -25,7 +25,7 @@
         SerializerLike,
         ShardsLike,
     )
-    from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
+    from zarr.core.array_spec import ArrayConfigLike
     from zarr.core.buffer import NDArrayLike
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
     from zarr.core.common import (
@@ -625,7 +625,7 @@ def create(
     codecs: Iterable[Codec | dict[str, JSON]] | None = None,
     dimension_names: Iterable[str] | None = None,
     storage_options: dict[str, Any] | None = None,
-    config: ArrayConfig | ArrayConfigLike | None = None,
+    config: ArrayConfigLike | None = None,
     **kwargs: Any,
 ) -> Array:
     """Create an array.
@@ -695,7 +695,7 @@ def create(
     storage_options : dict
         If using an fsspec URL to create the store, these will be passed to
         the backend implementation. Ignored otherwise.
-    config : ArrayConfig or ArrayConfigLike, optional
+    config : ArrayConfigLike, optional
         Runtime configuration of the array. If provided, will override the
         default values from `zarr.config.array`.
 
@@ -761,7 +761,7 @@ def create_array(
     dimension_names: Iterable[str] | None = None,
     storage_options: dict[str, Any] | None = None,
     overwrite: bool = False,
-    config: ArrayConfig | ArrayConfigLike | None = None,
+    config: ArrayConfigLike | None = None,
 ) -> Array:
     """Create an array.
 
@@ -853,7 +853,7 @@ def create_array(
         Ignored otherwise.
     overwrite : bool, default False
         Whether to overwrite an array with the same name in the store, if one exists.
-    config : ArrayConfig or ArrayConfigLike, optional
+    config : ArrayConfigLike, optional
         Runtime configuration for the array.
 
     Returns

src/zarr/core/array.py

@@ -219,7 +219,7 @@ class AsyncArray(Generic[T_ArrayMetadata]):
         The metadata of the array.
     store_path : StorePath
         The path to the Zarr store.
-    config : ArrayConfig, optional
+    config : ArrayConfigLike, optional
         The runtime configuration of the array, by default None.
 
     Attributes
@@ -244,22 +244,22 @@ def __init__(
         self: AsyncArray[ArrayV2Metadata],
         metadata: ArrayV2Metadata | ArrayV2MetadataDict,
         store_path: StorePath,
-        config: ArrayConfig | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> None: ...
 
     @overload
     def __init__(
         self: AsyncArray[ArrayV3Metadata],
         metadata: ArrayV3Metadata | ArrayV3MetadataDict,
         store_path: StorePath,
-        config: ArrayConfig | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> None: ...
 
     def __init__(
         self,
         metadata: ArrayMetadata | ArrayMetadataDict,
         store_path: StorePath,
-        config: ArrayConfig | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> None:
         if isinstance(metadata, dict):
             zarr_format = metadata["zarr_format"]
@@ -273,12 +273,11 @@ def __init__(
                 raise ValueError(f"Invalid zarr_format: {zarr_format}. Expected 2 or 3")
 
         metadata_parsed = parse_array_metadata(metadata)
-
-        config = ArrayConfig.from_dict({}) if config is None else config
+        config_parsed = parse_array_config(config)
 
         object.__setattr__(self, "metadata", metadata_parsed)
         object.__setattr__(self, "store_path", store_path)
-        object.__setattr__(self, "_config", config)
+        object.__setattr__(self, "_config", config_parsed)
         object.__setattr__(self, "codec_pipeline", create_codec_pipeline(metadata=metadata_parsed))
 
     # this overload defines the function signature when zarr_format is 2
@@ -302,7 +301,7 @@ async def create(
         # runtime
         overwrite: bool = False,
         data: npt.ArrayLike | None = None,
-        config: ArrayConfig | ArrayConfigLike | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> AsyncArray[ArrayV2Metadata]: ...
 
     # this overload defines the function signature when zarr_format is 3
@@ -331,7 +330,7 @@ async def create(
         # runtime
         overwrite: bool = False,
         data: npt.ArrayLike | None = None,
-        config: ArrayConfig | ArrayConfigLike | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> AsyncArray[ArrayV3Metadata]: ...
 
     @overload
@@ -359,7 +358,7 @@ async def create(
         # runtime
         overwrite: bool = False,
         data: npt.ArrayLike | None = None,
-        config: ArrayConfig | ArrayConfigLike | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> AsyncArray[ArrayV3Metadata]: ...
 
     @overload
@@ -393,7 +392,7 @@ async def create(
         # runtime
         overwrite: bool = False,
         data: npt.ArrayLike | None = None,
-        config: ArrayConfig | ArrayConfigLike | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]: ...
 
     @classmethod
@@ -428,7 +427,7 @@ async def create(
         # runtime
         overwrite: bool = False,
         data: npt.ArrayLike | None = None,
-        config: ArrayConfig | ArrayConfigLike | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
         """Method to create a new asynchronous array instance.
 
@@ -506,7 +505,7 @@ async def create(
             Whether to raise an error if the store already exists (default is False).
         data : npt.ArrayLike, optional
             The data to be inserted into the array (default is None).
-        config : ArrayConfig or ArrayConfigLike, optional
+        config : ArrayConfigLike, optional
             Runtime configuration for the array.
 
         Returns
@@ -569,7 +568,7 @@ async def _create(
         # runtime
         overwrite: bool = False,
         data: npt.ArrayLike | None = None,
-        config: ArrayConfig | ArrayConfigLike | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
         """Method to create a new asynchronous array instance.
         See :func:`AsyncArray.create` for more details.
@@ -1741,7 +1740,7 @@ def create(
         compressor: dict[str, JSON] | None = None,
         # runtime
         overwrite: bool = False,
-        config: ArrayConfig | ArrayConfigLike | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> Array:
         """Creates a new Array instance from an initialized store.
 
@@ -1870,7 +1869,7 @@ def _create(
         compressor: dict[str, JSON] | None = None,
         # runtime
         overwrite: bool = False,
-        config: ArrayConfig | ArrayConfigLike | None = None,
+        config: ArrayConfigLike | None = None,
     ) -> Array:
         """Creates a new Array instance from an initialized store.
         See :func:`Array.create` for more details.
@@ -3810,7 +3809,8 @@ async def init_array(
     chunk_key_encoding: ChunkKeyEncodingLike | None = None,
     dimension_names: Iterable[str] | None = None,
     overwrite: bool = False,
-) -> ArrayV3Metadata | ArrayV2Metadata:
+    config: ArrayConfigLike | None,
+) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]:
     """Create and persist an array metadata document.
 
     Parameters
@@ -3889,11 +3889,13 @@ async def init_array(
         Zarr format 3 only. Zarr format 2 arrays should not use this parameter.
     overwrite : bool, default False
         Whether to overwrite an array with the same name in the store, if one exists.
+    config : ArrayConfigLike or None, optional
+        Configuration for this array.
 
     Returns
     -------
-    ArrayV3Metadata | ArrayV2Metadata
-        The array metadata document.
+    AsyncArray
+        The AsyncArray.
     """
 
     if zarr_format is None:
@@ -3993,14 +3995,9 @@ async def init_array(
             attributes=attributes,
         )
 
-    # save the metadata to disk
-    # TODO: make this easier -- it should be a simple function call that takes a {key: buffer}
-    coros = (
-        (store_path / key).set(value)
-        for key, value in meta.to_buffer_dict(default_buffer_prototype()).items()
-    )
-    await gather(*coros)
-    return meta
+    arr = AsyncArray(metadata=meta, store_path=store_path, config=config)
+    await arr._save_metadata(meta, ensure_parents=True)
+    return arr
 
 
 async def create_array(
@@ -4023,7 +4020,7 @@ async def create_array(
     dimension_names: Iterable[str] | None = None,
     storage_options: dict[str, Any] | None = None,
     overwrite: bool = False,
-    config: ArrayConfig | ArrayConfigLike | None = None,
+    config: ArrayConfigLike | None = None,
     write_data: bool = True,
 ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array.
@@ -4113,7 +4110,7 @@ async def create_array(
         Ignored otherwise.
     overwrite : bool, default False
         Whether to overwrite an array with the same name in the store, if one exists.
-    config : ArrayConfig or ArrayConfigLike, optional
+    config : ArrayConfigLike, optional
         Runtime configuration for the array.
     write_data : bool
         If a pre-existing array-like object was provided to this function via the ``data`` parameter
@@ -4139,13 +4136,12 @@ async def create_array(
     <AsyncArray memory://140349042942400 shape=(100, 100) dtype=int32>
     """
     mode: Literal["a"] = "a"
-    config_parsed = parse_array_config(config)
     store_path = await make_store_path(store, path=name, mode=mode, storage_options=storage_options)
 
     data_parsed, shape_parsed, dtype_parsed = _parse_data_params(
         data=data, shape=shape, dtype=dtype
     )
-    meta = await init_array(
+    result = await init_array(
         store_path=store_path,
         shape=shape_parsed,
         dtype=dtype_parsed,
@@ -4161,9 +4157,9 @@ async def create_array(
         chunk_key_encoding=chunk_key_encoding,
         dimension_names=dimension_names,
         overwrite=overwrite,
+        config=config,
     )
 
-    result = AsyncArray(metadata=meta, store_path=store_path, config=config_parsed)
     if write_data is True and data_parsed is not None:
         await result._set_selection(
             BasicIndexer(..., shape=result.shape, chunk_grid=result.metadata.chunk_grid),

src/zarr/core/array_spec.py

@@ -21,7 +21,7 @@
     from zarr.core.common import ChunkCoords
 
 
-class ArrayConfigLike(TypedDict):
+class ArrayConfigParams(TypedDict):
     """
     A TypedDict model of the attributes of an ArrayConfig class, but with no required fields.
     This allows for partial construction of an ArrayConfig, with the assumption that the unset
@@ -56,13 +56,13 @@ def __init__(self, order: MemoryOrder, write_empty_chunks: bool) -> None:
         object.__setattr__(self, "write_empty_chunks", write_empty_chunks_parsed)
 
     @classmethod
-    def from_dict(cls, data: ArrayConfigLike) -> Self:
+    def from_dict(cls, data: ArrayConfigParams) -> Self:
         """
         Create an ArrayConfig from a dict. The keys of that dict are a subset of the
         attributes of the ArrayConfig class. Any keys missing from that dict will be set to the
         the values in the ``array`` namespace of ``zarr.config``.
         """
-        kwargs_out: ArrayConfigLike = {}
+        kwargs_out: ArrayConfigParams = {}
         for f in fields(ArrayConfig):
             field_name = cast(Literal["order", "write_empty_chunks"], f.name)
             if field_name not in data:
@@ -72,7 +72,10 @@ def from_dict(cls, data: ArrayConfigLike) -> Self:
         return cls(**kwargs_out)
 
 
-def parse_array_config(data: ArrayConfig | ArrayConfigLike | None) -> ArrayConfig:
+ArrayConfigLike = ArrayConfig | ArrayConfigParams
+
+
+def parse_array_config(data: ArrayConfigLike | None) -> ArrayConfig:
     """
     Convert various types of data to an ArrayConfig.
     """