fix: can now declare config in open_array

Author: d-v-b

src/zarr/api/asynchronous.py

@@ -1242,6 +1242,7 @@ async def open_array(
     zarr_format: ZarrFormat | None = None,
     path: PathLike = "",
     storage_options: dict[str, Any] | None = None,
+    config: ArrayConfigLike | None = None,
     **kwargs: Any,  # TODO: type kwargs as valid args to save
 ) -> AnyAsyncArray:
     """Open an array using file-mode-like semantics.
@@ -1261,6 +1262,8 @@ async def open_array(
     storage_options : dict
         If using an fsspec URL to create the store, these will be passed to
         the backend implementation. Ignored otherwise.
+    config : ArrayConfigLike
+        Declaration of the runtime configuration for the array.
     **kwargs
         Any keyword arguments to pass to [`create`][zarr.api.asynchronous.create].
 
@@ -1279,7 +1282,7 @@ async def open_array(
         _warn_write_empty_chunks_kwarg()
 
     try:
-        return await AsyncArray.open(store_path, zarr_format=zarr_format)
+        return await AsyncArray.open(store_path, zarr_format=zarr_format, config=config)
     except FileNotFoundError as err:
         if not store_path.read_only and mode in _CREATE_MODES:
             overwrite = _infer_overwrite(mode)

src/zarr/api/synchronous.py

@@ -1369,6 +1369,7 @@ def open_array(
     zarr_format: ZarrFormat | None = None,
     path: PathLike = "",
     storage_options: dict[str, Any] | None = None,
+    config: ArrayConfigLike | None = None,
     **kwargs: Any,
 ) -> AnyArray:
     """Open an array using file-mode-like semantics.
@@ -1388,6 +1389,8 @@ def open_array(
     storage_options : dict
         If using an fsspec URL to create the store, these will be passed to
         the backend implementation. Ignored otherwise.
+    config : ArrayConfigLike
+        Declaration of the runtime configuration for the array.
     **kwargs
         Any keyword arguments to pass to [`create`][zarr.api.asynchronous.create].
 
@@ -1405,6 +1408,7 @@ def open_array(
                 zarr_format=zarr_format,
                 path=path,
                 storage_options=storage_options,
+                config=config,
                 **kwargs,
             )
         )

src/zarr/codecs/sharding.py

@@ -348,9 +348,16 @@ def __getstate__(self) -> dict[str, Any]:
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         config = state["configuration"]
+        codec_class_map = parse_codec_class_map(None)
         object.__setattr__(self, "chunk_shape", parse_shapelike(config["chunk_shape"]))
-        object.__setattr__(self, "codecs", parse_codecs(config["codecs"]))
-        object.__setattr__(self, "index_codecs", parse_codecs(config["index_codecs"]))
+        object.__setattr__(
+            self, "codecs", parse_codecs(config["codecs"], codec_class_map=codec_class_map)
+        )
+        object.__setattr__(
+            self,
+            "index_codecs",
+            parse_codecs(config["index_codecs"], codec_class_map=codec_class_map),
+        )
         object.__setattr__(self, "index_location", parse_index_location(config["index_location"]))
 
         # Use instance-local lru_cache to avoid memory leaks

src/zarr/core/array.py

@@ -203,13 +203,13 @@ def _chunk_sizes_from_shape(
     return tuple(result)
 
 
-def parse_array_metadata(data: object, config: ArrayConfig) -> ArrayMetadata:
-    if isinstance(data, ArrayMetadata):
-        return data.with_config(config)
+def parse_array_metadata(data: object, codec_class_map: Mapping[str, type[Codec]]) -> ArrayMetadata:
+    if isinstance(data, ArrayV3Metadata):
+        return type(data).from_dict(data.to_dict(), codec_class_map=codec_class_map)
     elif isinstance(data, dict):
         zarr_format = data.get("zarr_format")
         if zarr_format == 3:
-            meta_out = ArrayV3Metadata.from_dict(data, config=config)
+            meta_out = ArrayV3Metadata.from_dict(data, codec_class_map=codec_class_map)
             if len(meta_out.storage_transformers) > 0:
                 msg = (
                     f"Array metadata contains storage transformers: {meta_out.storage_transformers}."
@@ -218,26 +218,37 @@ def parse_array_metadata(data: object, config: ArrayConfig) -> ArrayMetadata:
                 raise ValueError(msg)
             return meta_out
         elif zarr_format == 2:
-            return ArrayV2Metadata.from_dict(data, config=config)
+            return ArrayV2Metadata.from_dict(data)
         else:
             raise ValueError(f"Invalid zarr_format: {zarr_format}. Expected 2 or 3")
     raise TypeError  # pragma: no cover
 
 
-def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None) -> CodecPipeline:
+def create_codec_pipeline(
+    metadata: ArrayMetadata,
+    *,
+    store: Store | None = None,
+    config: ArrayConfig | None = None,
+) -> CodecPipeline:
+    pipeline_class: type[CodecPipeline]
+    if config is not None:
+        pipeline_class = config.codec_pipeline_class
+    else:
+        pipeline_class = get_pipeline_class()
+
     if store is not None:
         try:
-            return get_pipeline_class().from_array_metadata_and_store(
+            return pipeline_class.from_array_metadata_and_store(
                 array_metadata=metadata, store=store
             )
         except NotImplementedError:
             pass
 
     if isinstance(metadata, ArrayV3Metadata):
-        return get_pipeline_class().from_codecs(metadata.codecs)
+        return pipeline_class.from_codecs(metadata.codecs)
     elif isinstance(metadata, ArrayV2Metadata):
         v2_codec = V2Codec(filters=metadata.filters, compressor=metadata.compressor)
-        return get_pipeline_class().from_codecs([v2_codec])
+        return pipeline_class.from_codecs([v2_codec])
     raise TypeError  # pragma: no cover
 
 
@@ -360,7 +371,9 @@ def __init__(
         config: ArrayConfigLike | None = None,
     ) -> None:
         config_parsed = parse_array_config(config)
-        metadata_parsed = parse_array_metadata(metadata, config=config_parsed)
+        metadata_parsed = parse_array_metadata(
+            metadata, codec_class_map=config_parsed.codec_class_map
+        )
 
         object.__setattr__(self, "metadata", metadata_parsed)
         object.__setattr__(self, "store_path", store_path)
@@ -369,7 +382,9 @@ def __init__(
         object.__setattr__(
             self,
             "codec_pipeline",
-            create_codec_pipeline(metadata=metadata_parsed, store=store_path.store),
+            create_codec_pipeline(
+                metadata=metadata_parsed, store=store_path.store, config=config_parsed
+            ),
         )
 
     # this overload defines the function signature when zarr_format is 2
@@ -785,6 +800,7 @@ def _create_metadata_v3(
         codecs: Iterable[Codec | dict[str, JSON]] | None = None,
         dimension_names: DimensionNamesLike = None,
         attributes: dict[str, JSON] | None = None,
+        codec_class_map: Mapping[str, type[Codec]] | None = None,
     ) -> ArrayV3Metadata:
         """Create an instance of ArrayV3Metadata."""
         filters: tuple[ArrayArrayCodec, ...]
@@ -822,6 +838,7 @@ def _create_metadata_v3(
             codecs=codecs_parsed,  # type: ignore[arg-type]
             dimension_names=tuple(dimension_names) if dimension_names else None,
             attributes=attributes or {},
+            codec_class_map=codec_class_map,
         )
 
     @classmethod
@@ -869,6 +886,7 @@ async def _create_v3(
             codecs=codecs,
             dimension_names=dimension_names,
             attributes=attributes,
+            codec_class_map=config.codec_class_map,
         )
 
         array = cls(metadata=metadata, store_path=store_path, config=config)
@@ -993,14 +1011,18 @@ def from_dict(
         ValueError
             If the dictionary data is invalid or incompatible with either Zarr format 2 or 3 array creation.
         """
-        metadata = parse_array_metadata(data)
+        from zarr.core.array_spec import parse_codec_class_map
+
+        metadata = parse_array_metadata(data, codec_class_map=parse_codec_class_map(None))
         return cls(metadata=metadata, store_path=store_path)
 
     @classmethod
     async def open(
         cls,
         store: StoreLike,
         zarr_format: ZarrFormat | None = 3,
+        *,
+        config: ArrayConfigLike | None = None,
     ) -> AnyAsyncArray:
         """
         Async method to open an existing Zarr array from a given store.
@@ -1013,6 +1035,8 @@ async def open(
             for a description of all valid StoreLike values.
         zarr_format : ZarrFormat | None, optional
             The Zarr format version (default is 3).
+        config : ArrayConfigLike | None, (default is None)
+            Runtime configuration for the array.
 
         Returns
         -------
@@ -1044,7 +1068,7 @@ async def example():
         metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format)
         # TODO: remove this cast when we have better type hints
         _metadata_dict = cast("ArrayMetadataJSON_V3", metadata_dict)
-        return cls(store_path=store_path, metadata=_metadata_dict)
+        return cls(store_path=store_path, metadata=_metadata_dict, config=config)
 
     @property
     def store(self) -> Store:
@@ -4710,7 +4734,7 @@ async def init_array(
     chunk_key_encoding: ChunkKeyEncodingLike | None = None,
     dimension_names: DimensionNamesLike = None,
     overwrite: bool = False,
-    config: ArrayConfigLike | None = None,
+    config: ArrayConfig | None = None,
 ) -> AnyAsyncArray:
     """Create and persist an array metadata document.
 
@@ -4948,6 +4972,7 @@ async def init_array(
             codecs=codecs_out,
             dimension_names=dimension_names,
             attributes=attributes,
+            codec_class_map=config.codec_class_map if config is not None else None,
         )
 
     arr = AsyncArray(metadata=meta, store_path=store_path, config=config)
@@ -5145,7 +5170,7 @@ async def create_array(
             chunk_key_encoding=chunk_key_encoding,
             dimension_names=dimension_names,
             overwrite=overwrite,
-            config=config,
+            config=parse_array_config(config),
         )
 
 

src/zarr/core/array_spec.py

@@ -3,7 +3,7 @@
 from dataclasses import dataclass, fields
 from typing import TYPE_CHECKING, Any, Final, Literal, Self, TypedDict, cast
 
-from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec
+from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec, CodecPipeline
 from zarr.core.common import (
     MemoryOrder,
     parse_bool,
@@ -30,19 +30,19 @@ class CodecPipelineRequest(TypedDict):
     options: NotRequired[dict[str, object]]
 
 
-class ArrayConfigParams(TypedDict):
+class ArrayConfigParams(TypedDict, closed=True):  # type: ignore[call-arg]
     """
     A TypedDict model of the attributes of an ArrayConfig class.
     """
 
     order: MemoryOrder
     write_empty_chunks: bool
     read_missing_chunks: bool
-    codec_class_map: Mapping[str, object]
-    codec_pipeline_class: CodecPipelineRequest
+    codec_class_map: Mapping[str, type[ArrayArrayCodec | ArrayBytesCodec | BytesBytesCodec]]
+    codec_pipeline_class: type[CodecPipeline]
 
 
-class ArrayConfigRequest(TypedDict):
+class ArrayConfigRequest(TypedDict, closed=True):  # type: ignore[call-arg]
     """
     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
@@ -55,7 +55,7 @@ class ArrayConfigRequest(TypedDict):
     codec_class_map: NotRequired[
         Mapping[str, type[ArrayArrayCodec | ArrayBytesCodec | BytesBytesCodec]]
     ]
-    codec_pipeline_class: NotRequired[CodecPipelineRequest]
+    codec_pipeline_class: NotRequired[type[CodecPipeline]]
 
 
 ArrayConfigKeys = Literal[
@@ -112,16 +112,16 @@ class ArrayConfig:
     codec_class_map : Mapping[str, object]
         A codec name : codec class mapping that defines the codec classes available
         for array creation.
-    codec_pipeline_class : CodecPipelineRequest
-        A request for a pipeline class that will be used for orchestrating chunk encoding and
+    codec_pipeline_class : type[CodecPipeline]
+        A codec pipeline class that will be used for orchestrating chunk encoding and
         decoding.
     """
 
     order: MemoryOrder
     write_empty_chunks: bool
     read_missing_chunks: bool
     codec_class_map: Mapping[str, type[Codec]]
-    codec_pipeline_class: CodecPipelineRequest
+    codec_pipeline_class: type[CodecPipeline]
 
     def __init__(
         self,
@@ -131,7 +131,7 @@ def __init__(
         read_missing_chunks: bool = True,
         codec_class_map: Mapping[str, type[ArrayBytesCodec | ArrayArrayCodec | BytesBytesCodec]]
         | None = None,
-        codec_pipeline_class: CodecPipelineRequest | None = None,
+        codec_pipeline_class: type[CodecPipeline] | None = None,
     ) -> None:
         order_parsed = parse_order(order)
         write_empty_chunks_parsed = parse_bool(write_empty_chunks)
@@ -213,19 +213,16 @@ def _import_by_name(path: str) -> object | type:
     return obj
 
 
-def parse_codec_pipeline_class(obj: CodecPipelineRequest | None) -> CodecPipelineRequest:
+def parse_codec_pipeline_class(obj: type[CodecPipeline] | None) -> type[CodecPipeline]:
     if obj is None:
-        config_entry: dict[str, str | int] = zarr_config.get("codec_pipeline")
+        config_entry: dict[str, str] = zarr_config.get("codec_pipeline")
         if "path" not in config_entry:
             msg = (
                 "The codec_pipeline field in the global config is malformed. "
                 "Expected 'path' key was not found."
             )
             raise KeyError(msg)
-        else:
-            path = config_entry["path"]
-        options = {"batch_size": config_entry.get("batch_size", 1)}
-        return {"class_path": path, "options": options}
+        return _import_by_name(config_entry["path"])  # type: ignore[return-value]
     return obj
 
 
@@ -241,6 +238,9 @@ def parse_codec_class_map(obj: Mapping[str, type[Codec]] | None) -> Mapping[str,
         out: dict[str, type[Codec]] = {}
         for key, value in name_map.items():
             maybe_cls = _import_by_name(value)
+            if not isinstance(maybe_cls, type):
+                msg = f"Expected a type, got {maybe_cls}"
+                raise TypeError(msg)
             if not issubclass(maybe_cls, Codec):
                 msg = f"Expected a subclass of `Codec`, got {maybe_cls}"
                 raise TypeError(msg)