feat(experimental): scope ArrayIO's extension resolution with a Context

Extends zarr.core.context.Context — the extensions-in-scope value the
structured data type resolution already uses — with codecs and
codec_pipeline, and threads an optional context through the resolution
functions: parse_dtype and the data-type getters, codec normalization
(_resolve_codec, _parse_*_codec, parse_codecs), metadata from_dict on
both formats, parse_array_metadata, create_codec_pipeline, and
build_array_metadata. A kind the context leaves None resolves from the
default registries exactly as before; a kind it sets is a closed world,
and a name outside it fails to resolve, naming the scope.

AsyncArrayIO and ArrayIO take the context in build, open and their
constructors, making them the first public takers, and AsyncArray
passes one through to the AsyncArrayIO it wraps.

Refs zarr-developers#4217

Assisted-by: ClaudeCode:claude-fable-5

Author: d-v-b

changes/420.feature.md

@@ -0,0 +1 @@
+Added a `context` to `zarr.experimental.ArrayIO` and `zarr.experimental.AsyncArrayIO`: a value defining the data types, codecs and codec pipeline in scope for one array. Extension names in a stored document or in creation keyword arguments resolve through the context instead of the default registries; a kind the context leaves unset falls back to the defaults, while a kind it sets is a closed world, and a name outside it fails to resolve with an error naming the scope.

src/zarr/core/array.py

@@ -170,6 +170,7 @@
     from zarr.abc.store import Store
     from zarr.codecs.sharding import IndexLocation, ShardingCodec
     from zarr.core.buffer import Buffer
+    from zarr.core.context import Context
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar
     from zarr.storage import StoreLike
     from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3
@@ -215,9 +216,13 @@ def _chunk_sizes_from_shape(
     return tuple(result)
 
 
-def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata:
+def parse_array_metadata(
+    data: Any, path: str | None = None, *, context: Context | None = None
+) -> ArrayMetadata:
     """Array metadata from a metadata object or a metadata document, naming the array at
-    `path` in warnings about how an invalid document was read.
+    `path` in warnings about how an invalid document was read. With a `context`, a
+    document's extension names resolve through it instead of the default registries;
+    a metadata object is already resolved and passes through any context.
 
     `ArrayV2Metadata` accepts a chunk size of 0, as it always has, though only an
     invalid document holds one: such metadata is read as the documents it would store
@@ -230,7 +235,7 @@ def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata:
     if isinstance(data, dict):
         zarr_format = data.get("zarr_format")
         if zarr_format == 3:
-            meta_out = ArrayV3Metadata.from_dict(data, path=path)
+            meta_out = ArrayV3Metadata.from_dict(data, path=path, context=context)
             if len(meta_out.storage_transformers) > 0:
                 msg = (
                     f"Array metadata contains storage transformers: {meta_out.storage_transformers}."
@@ -239,16 +244,19 @@ def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata:
                 raise ValueError(msg)
             return meta_out
         elif zarr_format == 2:
-            return ArrayV2Metadata.from_dict(data, path=path)
+            return ArrayV2Metadata.from_dict(data, path=path, context=context)
         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, context: Context | None = None
+) -> CodecPipeline:
+    pipeline_class = get_pipeline_class() if context is None else context.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:
@@ -261,7 +269,7 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None
         # "disables partial reads" warning) fire here; `evolve_from_array_spec`
         # re-splits the same already-warned-about chain via
         # `codecs_from_list_unchecked`, so it does not re-emit them.
-        pipeline = get_pipeline_class().from_codecs(metadata.codecs)
+        pipeline = pipeline_class.from_codecs(metadata.codecs)
 
         # Use the regular chunk shape if available, otherwise use a
         # placeholder. The ChunkTransform is shape-agnostic — the actual
@@ -280,7 +288,7 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None
         return pipeline.evolve_from_array_spec(chunk_spec)
     elif isinstance(metadata, ArrayV2Metadata):
         v2_codec = V2Codec(filters=metadata.filters, compressor=metadata.compressor)
-        pipeline = get_pipeline_class().from_codecs([v2_codec])
+        pipeline = pipeline_class.from_codecs([v2_codec])
         chunk_spec = ArraySpec(
             shape=metadata.chunks,
             dtype=metadata.dtype,
@@ -403,6 +411,7 @@ def __init__(
         metadata: ArrayV2Metadata | ArrayV2MetadataDict,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
+        context: Context | None = None,
     ) -> None: ...
 
     @overload
@@ -411,15 +420,17 @@ def __init__(
         metadata: ArrayV3Metadata | ArrayMetadataJSON_V3,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
+        context: Context | None = None,
     ) -> None: ...
 
     def __init__(
         self,
         metadata: ArrayMetadata | ArrayMetadataDict,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
+        context: Context | None = None,
     ) -> None:
-        object.__setattr__(self, "_io", AsyncArrayIO(metadata, store_path, config))
+        object.__setattr__(self, "_io", AsyncArrayIO(metadata, store_path, config, context))
 
     @classmethod
     def _from_io(cls, io: AsyncArrayIO) -> AnyAsyncArray:
@@ -4599,21 +4610,25 @@ def build_array_metadata(
     attributes: dict[str, JSON] | None = None,
     chunk_key_encoding: ChunkKeyEncodingLike | None = None,
     dimension_names: DimensionNamesLike = None,
+    context: Context | None = None,
 ) -> ArrayV2Metadata | ArrayV3Metadata:
     """Build an array metadata document from creation-style keyword arguments.
 
     This is the pure half of `init_array`: it normalizes the same keyword
     arguments into an in-memory metadata object without touching any store.
     The parameters have the same meaning as in
-    [`zarr.api.asynchronous.create_array`][].
+    [`zarr.api.asynchronous.create_array`][]. With a `context`, the data type
+    and Zarr format 3 codec names resolve through it instead of the default
+    registries; Zarr format 2 filters and compressors are numcodecs objects,
+    which no context governs.
     """
 
     if zarr_format is None:
         zarr_format = _default_zarr_format()
 
     from zarr.codecs.sharding import ShardingCodec
 
-    zdtype = parse_dtype(dtype, zarr_format=zarr_format)
+    zdtype = parse_dtype(dtype, zarr_format=zarr_format, context=context)
     shape_parsed = parse_shapelike(shape)
     item_size = zdtype.item_size if isinstance(zdtype, HasItemSize) else 1
     chunk_key_encoding_parsed = _parse_chunk_key_encoding(
@@ -4689,6 +4704,7 @@ def build_array_metadata(
             filters=filters,
             serializer=serializer,
             dtype=zdtype,
+            context=context,
         )
         sub_codecs = cast("tuple[Codec, ...]", (*array_array, array_bytes, *bytes_bytes))
         # A stored chunk grid passed as chunks= / shards= (e.g. by the "keep"
@@ -5312,6 +5328,7 @@ def _parse_chunk_encoding_v3(
     filters: FiltersLike,
     serializer: SerializerLike,
     dtype: ZDType[TBaseDType, TBaseScalar],
+    context: Context | None = None,
 ) -> tuple[tuple[ArrayArrayCodec, ...], ArrayBytesCodec, tuple[BytesBytesCodec, ...]]:
     """
     Generate chunk encoding classes for v3 arrays with optional defaults.
@@ -5328,15 +5345,17 @@ def _parse_chunk_encoding_v3(
             maybe_array_array = (filters,)
         else:
             maybe_array_array = cast("Iterable[Codec | dict[str, JSON]]", filters)
-        out_array_array = tuple(_parse_array_array_codec(c) for c in maybe_array_array)
+        out_array_array = tuple(
+            _parse_array_array_codec(c, context=context) for c in maybe_array_array
+        )
 
     if serializer == "auto":
         out_array_bytes = default_serializer_v3(dtype)
     else:
         # TODO: ensure that the serializer is compatible with the ndarray produced by the
         # array-array codecs. For example, if a sequence of array-array codecs produces an
         # array with a single-byte data type, then the serializer should not specify endiannesss.
-        out_array_bytes = _parse_array_bytes_codec(serializer)
+        out_array_bytes = _parse_array_bytes_codec(serializer, context=context)
 
     if compressors is None:
         out_bytes_bytes: tuple[BytesBytesCodec, ...] = ()
@@ -5349,7 +5368,9 @@ def _parse_chunk_encoding_v3(
         else:
             maybe_bytes_bytes = cast("Iterable[Codec | dict[str, JSON]]", compressors)
 
-        out_bytes_bytes = tuple(_parse_bytes_bytes_codec(c) for c in maybe_bytes_bytes)
+        out_bytes_bytes = tuple(
+            _parse_bytes_bytes_codec(c, context=context) for c in maybe_bytes_bytes
+        )
 
     # TODO: ensure that the serializer is compatible with the ndarray produced by the
     # array-array codecs. For example, if a sequence of array-array codecs produces an
@@ -6207,11 +6228,18 @@ class AsyncArrayIO:
     opened from a stored document that needed a repair stores that repair
     before its first chunk write. Call `store_metadata` first to make that
     write explicit.
+
+    A `context` (see [`zarr.core.context.Context`][]) defines the data types,
+    codecs and codec pipeline in scope for this array: extension names in a
+    stored document or in `build`'s keyword arguments resolve through it, and
+    a name outside it fails to resolve even when the default registries know
+    it. Without one, the default registries are used.
     """
 
     _metadata: ArrayMetadata
     _store_path: StorePath
     _config: ArrayConfig
+    _context: Context | None
     _codec_pipeline: CodecPipeline
     _chunk_grid: ChunkGrid
 
@@ -6220,14 +6248,16 @@ def __init__(
         metadata: ArrayMetadata | ArrayMetadataDict,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
+        context: Context | None = None,
     ) -> None:
-        metadata_parsed = parse_array_metadata(metadata, str(store_path))
+        metadata_parsed = parse_array_metadata(metadata, str(store_path), context=context)
         self._metadata = metadata_parsed
         self._store_path = store_path
         self._config = parse_array_config(config)
+        self._context = context
         self._chunk_grid = ChunkGrid.from_metadata(metadata_parsed)
         self._codec_pipeline = create_codec_pipeline(
-            metadata=metadata_parsed, store=store_path.store
+            metadata=metadata_parsed, store=store_path.store, context=context
         )
 
     @classmethod
@@ -6250,12 +6280,15 @@ async def build(
         chunk_key_encoding: ChunkKeyEncodingLike | None = None,
         dimension_names: DimensionNamesLike = None,
         config: ArrayConfigLike | None = None,
+        context: Context | None = None,
     ) -> AsyncArrayIO:
         """Declare an array bound to a location in a store, without storing anything.
 
         The keyword arguments have the same meaning as in
         [`zarr.api.asynchronous.create_array`][]. The returned instance holds
-        the normalized metadata in memory; `store_metadata` stores it.
+        the normalized metadata in memory; `store_metadata` stores it. With a
+        `context`, the data type and codec names resolve through it instead of
+        the default registries.
         """
         metadata = build_array_metadata(
             shape=shape,
@@ -6271,9 +6304,10 @@ async def build(
             attributes=attributes,
             chunk_key_encoding=chunk_key_encoding,
             dimension_names=dimension_names,
+            context=context,
         )
         store_path = await make_store_path(store, path=name)
-        return cls(metadata, store_path, config)
+        return cls(metadata, store_path, config, context=context)
 
     @classmethod
     async def open(
@@ -6283,19 +6317,24 @@ async def open(
         name: str | None = None,
         zarr_format: ZarrFormat | None = None,
         config: ArrayConfigLike | None = None,
+        context: Context | None = None,
     ) -> AsyncArrayIO:
         """Read the metadata of an existing array and bind to it.
 
         This is the one classmethod constructor that performs IO: the array's
         stored metadata document is read from the store. `zarr_format=None`
         (the default) reads Zarr format 3 and falls back to format 2, as
-        `open` does in `zarrs`; pass 2 or 3 to require one format.
+        `open` does in `zarrs`; pass 2 or 3 to require one format. With a
+        `context`, the document's data type and codec names resolve through it
+        instead of the default registries.
         """
         store_path = await make_store_path(store, path=name)
         metadata = parse_array_metadata(
-            await get_array_metadata(store_path, zarr_format=zarr_format), str(store_path)
+            await get_array_metadata(store_path, zarr_format=zarr_format),
+            str(store_path),
+            context=context,
         )
-        return cls(metadata, store_path, config)
+        return cls(metadata, store_path, config, context=context)
 
     @property
     def metadata(self) -> ArrayMetadata:
@@ -6350,7 +6389,9 @@ async def verify(self) -> None:
         """
         bound = self._metadata
         stored = parse_array_metadata(
-            await get_array_metadata(self._store_path, zarr_format=None), str(self._store_path)
+            await get_array_metadata(self._store_path, zarr_format=None),
+            str(self._store_path),
+            context=self._context,
         )
         if stored == bound:
             return
@@ -6515,8 +6556,9 @@ def __init__(
         metadata: ArrayMetadata | ArrayMetadataDict,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
+        context: Context | None = None,
     ) -> None:
-        self._async_io = AsyncArrayIO(metadata, store_path, config)
+        self._async_io = AsyncArrayIO(metadata, store_path, config, context)
 
     @classmethod
     def _wrap(cls, async_io: AsyncArrayIO) -> ArrayIO:
@@ -6544,6 +6586,7 @@ def build(
         chunk_key_encoding: ChunkKeyEncodingLike | None = None,
         dimension_names: DimensionNamesLike = None,
         config: ArrayConfigLike | None = None,
+        context: Context | None = None,
     ) -> ArrayIO:
         """Declare an array bound to a location in a store, without storing anything.
 
@@ -6568,6 +6611,7 @@ def build(
                     chunk_key_encoding=chunk_key_encoding,
                     dimension_names=dimension_names,
                     config=config,
+                    context=context,
                 )
             )
         )
@@ -6580,13 +6624,18 @@ def open(
         name: str | None = None,
         zarr_format: ZarrFormat | None = None,
         config: ArrayConfigLike | None = None,
+        context: Context | None = None,
     ) -> ArrayIO:
         """Read the metadata of an existing array and bind to it.
 
         See `AsyncArrayIO.open`.
         """
         return cls._wrap(
-            sync(AsyncArrayIO.open(store, name=name, zarr_format=zarr_format, config=config))
+            sync(
+                AsyncArrayIO.open(
+                    store, name=name, zarr_format=zarr_format, config=config, context=context
+                )
+            )
         )
 
     @property