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

changes/420.feature.md

@@ -0,0 +1,4 @@
+Added a `context` to `zarr.ArrayIO` and `zarr.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.
+`with_context` rebinds one array under another scope, so tests and applications select a codec pipeline per binding instead of through the global configuration; the pipeline-parity tests now run that way.
+Without a context, every binding uses the default registries and the globally configured pipeline, exactly as before: the context is the explicit, per-binding replacement for that global state, and the path to deprecating it.
+The scope governs top-level names: codecs nested inside a `sharding_indexed` codec resolve through the default registry today. `with_context` re-resolves metadata that was already resolved, so a stored name outside the new scope fails there too.

src/zarr/core/context.py

@@ -7,15 +7,22 @@
 threaded through the read by zarr. (The pair mirrors `referencing.Registry` and
 `referencing.Resolver` in python-jsonschema.)
 
-Both are internal: nothing in the public API takes a context yet.
+The explicit-IO array layer (`zarr.ArrayIO` and `zarr.AsyncArrayIO`) takes a context publicly;
+everywhere else, both remain internal.
 """
 
 from __future__ import annotations
 
 from dataclasses import dataclass, replace
+from types import MappingProxyType
 from typing import TYPE_CHECKING
 
+from zarr.errors import UnknownCodecError
+
 if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from zarr.abc.codec import Codec, CodecPipeline
     from zarr.core.common import ZarrFormat
     from zarr.core.dtype.common import DTypeJSON
     from zarr.core.dtype.registry import DataTypeRegistry
@@ -27,13 +34,47 @@ class Context:
     """
     The extensions that Zarr metadata is read with, by kind.
 
+    A kind that is `None` resolves from the default registries, exactly as it
+    would with no context; a kind that is set is a closed world, and a name
+    outside it fails to resolve even when the default registries know it.
+    The scope governs top-level names: codecs nested inside a
+    `sharding_indexed` codec resolve through the default registry today.
+
     Attributes
     ----------
     data_types : DataTypeRegistry
         The data types a read may resolve.
+    codecs : Mapping[str, type[Codec]] | None
+        The codec class each codec name resolves to, or `None` for the
+        default codec registry. Codec instances are never resolved by name,
+        so they pass through any context.
+    codec_pipeline : type[CodecPipeline] | None
+        The codec pipeline class arrays are read and written with, or `None`
+        for the pipeline the global configuration selects.
     """
 
     data_types: DataTypeRegistry
+    codecs: Mapping[str, type[Codec]] | None = None
+    codec_pipeline: type[CodecPipeline] | None = None
+
+    def __post_init__(self) -> None:
+        # A read-only copy: the scope must not change under a binding that
+        # captured it. Contexts compare by their fields and are not hashable.
+        if self.codecs is not None:
+            object.__setattr__(self, "codecs", MappingProxyType(dict(self.codecs)))
+
+    def __getstate__(self) -> dict[str, object]:
+        # A mapping proxy does not pickle; carry the scope as a plain dict.
+        state = dict(self.__dict__)
+        if self.codecs is not None:
+            state["codecs"] = dict(self.codecs)
+        return state
+
+    def __setstate__(self, state: dict[str, object]) -> None:
+        self.__dict__.update(state)
+        codecs = state.get("codecs")
+        if codecs is not None:
+            object.__setattr__(self, "codecs", MappingProxyType(dict(codecs)))  # type: ignore[call-overload]
 
     @classmethod
     def default(cls) -> Context:
@@ -48,6 +89,31 @@ def default(cls) -> Context:
 
         return cls(data_types=data_type_registry)
 
+    def codec_class(self, name: str) -> type[Codec]:
+        """The codec class `name` resolves to in this context.
+
+        Raises [`zarr.errors.UnknownCodecError`][] naming the codecs in scope
+        when `name` is not among them.
+        """
+        if self.codecs is None:
+            from zarr.registry import get_codec_class
+
+            return get_codec_class(name)
+        try:
+            return self.codecs[name]
+        except KeyError:
+            raise UnknownCodecError(
+                f"Codec {name!r} is not in this context. Codecs in scope: {sorted(self.codecs)}."
+            ) from None
+
+    def pipeline_class(self) -> type[CodecPipeline]:
+        """The codec pipeline class this context resolves to."""
+        if self.codec_pipeline is None:
+            from zarr.registry import get_pipeline_class
+
+            return get_pipeline_class()
+        return self.codec_pipeline
+
 
 type JSONLocation = tuple[str | int, ...]
 

src/zarr/core/dtype/__init__.py

@@ -35,6 +35,7 @@
 
 if TYPE_CHECKING:
     from zarr.core.common import ZarrFormat
+    from zarr.core.context import Context
     from zarr.core.dtype.common import DTypeJSON
 
 from collections.abc import Mapping
@@ -161,9 +162,14 @@
 
 
 # TODO: find a better name for this function
-def get_data_type_from_native_dtype(dtype: npt.DTypeLike) -> ZDType[TBaseDType, TBaseScalar]:
+def get_data_type_from_native_dtype(
+    dtype: npt.DTypeLike, *, context: Context | None = None
+) -> ZDType[TBaseDType, TBaseScalar]:
     """
     Get a data type wrapper (an instance of ``ZDType``) from a native data type, e.g. a numpy dtype.
+
+    With a `context`, the match is made against the data types it holds instead of the
+    default data type registry.
     """
     if not isinstance(dtype, np.dtype):
         na_dtype: np.dtype[np.generic]
@@ -174,17 +180,22 @@ def get_data_type_from_native_dtype(dtype: npt.DTypeLike) -> ZDType[TBaseDType,
             na_dtype = np.dtype(dtype)
     else:
         na_dtype = dtype
-    return data_type_registry.match_dtype(dtype=na_dtype)
+    registry = data_type_registry if context is None else context.data_types
+    return registry.match_dtype(dtype=na_dtype)
 
 
 def get_data_type_from_json(
-    dtype_spec: DTypeJSON, *, zarr_format: ZarrFormat
+    dtype_spec: DTypeJSON, *, zarr_format: ZarrFormat, context: Context | None = None
 ) -> ZDType[TBaseDType, TBaseScalar]:
     """
     Given a JSON representation of a data type and a Zarr format version,
     attempt to create a ZDType instance from the registered ZDType classes.
+
+    With a `context`, the match is made against the data types it holds instead of the
+    default data type registry.
     """
-    return data_type_registry.match_json(dtype_spec, zarr_format=zarr_format)
+    registry = data_type_registry if context is None else context.data_types
+    return registry.match_json(dtype_spec, zarr_format=zarr_format)
 
 
 def parse_data_type(
@@ -234,6 +245,7 @@ def parse_dtype(
     dtype_spec: ZDTypeLike,
     *,
     zarr_format: ZarrFormat,
+    context: Context | None = None,
 ) -> ZDType[TBaseDType, TBaseScalar]:
     """
     Convert the input as a ZDType.
@@ -272,17 +284,24 @@ def parse_dtype(
     # First attempt to interpret the input as JSON
     if isinstance(dtype_spec, Mapping | str | Sequence):
         try:
-            return get_data_type_from_json(dtype_spec, zarr_format=zarr_format)
+            return get_data_type_from_json(dtype_spec, zarr_format=zarr_format, context=context)
         except ValueError:
             # no data type matched this JSON-like input
             pass
     if dtype_spec in VLEN_UTF8_ALIAS:
         # If the dtype request is one of the aliases for variable-length UTF-8 strings,
         # return that dtype.
-        return VariableLengthUTF8()  # type: ignore[return-value]
+        utf8 = VariableLengthUTF8()
+        if context is not None:
+            # The shortcut must not widen a closed world, and the context may
+            # map this JSON to its own class: return the registry's match.
+            return context.data_types.match_json(
+                utf8.to_json(zarr_format=zarr_format), zarr_format=zarr_format
+            )
+        return utf8  # type: ignore[return-value]
     # otherwise, we have either a numpy dtype string, or a zarr v3 dtype string, and in either case
     # we can create a native dtype from it, and do the dtype inference from that
-    return get_data_type_from_native_dtype(dtype_spec)  # type: ignore[arg-type]
+    return get_data_type_from_native_dtype(dtype_spec, context=context)  # type: ignore[arg-type]
 
 
 def __getattr__(name: str) -> object:

src/zarr/core/metadata/v2.py

@@ -20,6 +20,7 @@
 
     from zarr.core.buffer import Buffer, BufferPrototype
     from zarr.core.chunk_grids import ChunkGrid
+    from zarr.core.context import Context
     from zarr.core.dtype.wrapper import (
         TBaseDType,
         TBaseScalar,
@@ -152,10 +153,14 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
         }
 
     @classmethod
-    def from_dict(cls, data: dict[str, Any], *, path: str | None = None) -> ArrayV2Metadata:
+    def from_dict(
+        cls, data: dict[str, Any], *, path: str | None = None, context: Context | None = None
+    ) -> ArrayV2Metadata:
         """Read a stored `.zarray` document (with its attributes). An invalid document
         that `zarr.core.metadata.repair` can read is read as repaired; a reading the user
-        must act on warns, naming the array at `path`."""
+        must act on warns, naming the array at `path`. With a `context`, the document's
+        data type resolves through it instead of the default registry; filters and the
+        compressor are numcodecs objects, which no context governs."""
         repaired, readings = repair_array_document(data, 2)
         # a new dict, because we are modifying it
         _data: dict[str, Any] = dict(repaired)
@@ -178,7 +183,7 @@ def from_dict(cls, data: dict[str, Any], *, path: str | None = None) -> ArrayV2M
             "name": _data["dtype"],
             "object_codec_id": object_codec_id,
         }
-        dtype = get_data_type_from_json(dtype_spec, zarr_format=2)
+        dtype = get_data_type_from_json(dtype_spec, zarr_format=2, context=context)
 
         _data["dtype"] = dtype
         fill_value_encoded = _data.get("fill_value")

src/zarr/core/metadata/v3.py

@@ -53,6 +53,7 @@
     from zarr.codecs.sharding import ShardingCodec
     from zarr.core.buffer import Buffer, BufferPrototype
     from zarr.core.chunk_grids import ChunkGrid
+    from zarr.core.context import Context
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar
     from zarr.core.metadata.repair import ArrayDocument
 
@@ -70,7 +71,7 @@ def parse_node_type_array(data: object) -> Literal["array"]:
     )
 
 
-def parse_codecs(data: object) -> tuple[Codec, ...]:
+def parse_codecs(data: object, *, context: Context | None = None) -> tuple[Codec, ...]:
     out: tuple[Codec, ...] = ()
 
     if not isinstance(data, Iterable):
@@ -84,7 +85,11 @@ def parse_codecs(data: object) -> tuple[Codec, ...]:
         else:
             name_parsed, _ = parse_named_configuration(c, require_configuration=False)
 
-            codec_cls = get_codec_class(name_parsed)
+            codec_cls = (
+                get_codec_class(name_parsed)
+                if context is None
+                else context.codec_class(name_parsed)
+            )
             try:
                 out += (codec_cls.from_dict(c),)
             except KeyError as e:
@@ -665,11 +670,14 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
         return {ZARR_JSON: json_to_buffer(self.to_dict(), prototype=prototype, indent=indent)}
 
     @classmethod
-    def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self:
+    def from_dict(
+        cls, data: dict[str, JSON], *, path: str | None = None, context: Context | None = None
+    ) -> Self:
         """Read a stored `zarr.json` array document. An invalid document that
         `zarr.core.metadata.repair` can read is read as repaired; a reading the user
         must act on warns, and a document the rectilinear chunks flag refuses raises,
-        naming the array at `path`."""
+        naming the array at `path`. With a `context`, the document's data type and
+        codec names resolve through it instead of the default registries."""
         # The flag gates what the document declares, so it is checked before repairs.
         chunk_grid = data.get("chunk_grid")
         if isinstance(chunk_grid, Mapping) and chunk_grid.get("name") == "rectilinear":
@@ -691,7 +699,7 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self:
         data_type_json = _data.pop("data_type")
         if not check_dtype_spec_v3(data_type_json):
             raise ValueError(f"Invalid data_type: {data_type_json!r}")
-        data_type = get_data_type_from_json(data_type_json, zarr_format=3)
+        data_type = get_data_type_from_json(data_type_json, zarr_format=3, context=context)
 
         # check that the fill value is consistent with the data type
         try:
@@ -713,7 +721,9 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self:
             shape=_data_typed["shape"],
             chunk_grid=_data_typed["chunk_grid"],  # type: ignore[arg-type]
             chunk_key_encoding=_data_typed["chunk_key_encoding"],  # type: ignore[arg-type]
-            codecs=_data_typed["codecs"],
+            # Codec names resolve here, under the context; `__init__`'s own
+            # `parse_codecs` passes the resolved instances through untouched.
+            codecs=parse_codecs(_data_typed["codecs"], context=context),
             # Attribute values are arbitrary JSON, so they have no field-specific
             # schema to validate. `__init__` checks the outer dict via `parse_attributes`.
             attributes=_data_typed.get("attributes", {}),  # type: ignore[arg-type]

src/zarr/experimental/__init__.py

@@ -1,5 +1,6 @@
 """The experimental module is a site for exporting new or experimental Zarr features."""
 
 from zarr.core.chunk_grids import ChunkGrid, ChunkSpec
+from zarr.core.context import Context
 
-__all__ = ["ChunkGrid", "ChunkSpec"]
+__all__ = ["ChunkGrid", "ChunkSpec", "Context"]

tests/test_array.py

@@ -56,6 +56,7 @@
     resolve_outer_and_inner_chunks,
 )
 from zarr.core.chunk_key_encodings import ChunkKeyEncodingParams
+from zarr.core.codec_pipeline import FusedCodecPipeline
 from zarr.core.common import JSON, ChunksLike, ZarrFormat, ceildiv
 from zarr.core.dtype import (
     DateTime64,
@@ -68,6 +69,7 @@
     VariableLengthBytes,
     VariableLengthUTF8,
     ZDType,
+    data_type_registry,
     parse_dtype,
 )
 from zarr.core.dtype.common import ENDIANNESS_STR, EndiannessStr
@@ -81,6 +83,7 @@
     ContainsGroupError,
     ZarrUserWarning,
 )
+from zarr.experimental import Context
 from zarr.storage import LocalStore, MemoryStore, StorePath
 from zarr.storage._logging import LoggingStore
 from zarr.types import AnyArray, AnyAsyncArray
@@ -2171,20 +2174,22 @@ def test_multiprocessing_fork_codec_pipeline_pool(store: Store) -> None:
     pool entirely.
     """
     data = np.arange(100).reshape(10, 10)
-    with zarr.config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}):
-        arr = zarr.create_array(store=store, shape=data.shape, chunks=(5, 5), dtype=data.dtype)
-        arr[:] = data
-        # Two multi-chunk reads: spawn 4 pool threads and leave all of their
-        # idle-semaphore permits released before the fork.
-        arr[:]
-        arr[:]
-        ctx = mp.get_context("fork")
-        with ctx.Pool() as pool:
-            # [:5] spans 2 chunks — below the ~4 inherited idle permits, so on
-            # unfixed code both submissions are swallowed by the dead pool.
-            result = pool.starmap_async(_index_array, [(arr, np.s_[:5])])
-            (read_back,) = result.get(timeout=30)
-        assert np.array_equal(read_back, data[:5])
+    fused = Context(data_types=data_type_registry, codec_pipeline=FusedCodecPipeline)
+    arr = zarr.create_array(store=store, shape=data.shape, chunks=(5, 5), dtype=data.dtype)
+    # The pipeline is chosen per binding, not through the global config.
+    arr = arr.io.with_context(fused).to_array()
+    arr[:] = data
+    # Two multi-chunk reads: spawn 4 pool threads and leave all of their
+    # idle-semaphore permits released before the fork.
+    arr[:]
+    arr[:]
+    ctx = mp.get_context("fork")
+    with ctx.Pool() as pool:
+        # [:5] spans 2 chunks — below the ~4 inherited idle permits, so on
+        # unfixed code both submissions are swallowed by the dead pool.
+        result = pool.starmap_async(_index_array, [(arr, np.s_[:5])])
+        (read_back,) = result.get(timeout=30)
+    assert np.array_equal(read_back, data[:5])
 
 
 def test_create_array_method_signature() -> None: