From 6e914021e25fdcae26772df0957eae5b04063f20 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 29 Oct 2026 14:21:04 +0100 Subject: [PATCH] feat(context): scope ArrayIO's extension resolution with a Context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. ArrayIO and AsyncArrayIO take the context in build, open and their constructors, and AsyncArray passes one through to the AsyncArrayIO it wraps. With the document-first binding the context is consumed only where the document resolves into classes: opening an array whose codec names fall outside the context still works for document operations, and resolution fails when the data path is built. Refs zarr-developers/zarr-python#4217 Assisted-by: ClaudeCode:claude-fable-5 --- changes/420.feature.md | 1 + src/zarr/core/array.py | 90 +++++++++++++++++++++++-------- src/zarr/core/context.py | 46 +++++++++++++++- src/zarr/core/dtype/__init__.py | 24 ++++++--- src/zarr/core/metadata/v2.py | 11 ++-- src/zarr/core/metadata/v3.py | 22 +++++--- src/zarr/experimental/__init__.py | 3 +- src/zarr/registry.py | 28 +++++++--- tests/test_array_io.py | 75 +++++++++++++++++++++++++- 9 files changed, 252 insertions(+), 48 deletions(-) create mode 100644 changes/420.feature.md diff --git a/changes/420.feature.md b/changes/420.feature.md new file mode 100644 index 0000000000..e9f5e42c8a --- /dev/null +++ b/changes/420.feature.md @@ -0,0 +1 @@ +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. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 33b7a7daef..8d6f6d00cc 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -175,6 +175,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 @@ -220,9 +221,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 @@ -235,7 +240,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}." @@ -244,16 +249,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: @@ -266,7 +274,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 @@ -285,7 +293,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, @@ -408,6 +416,7 @@ def __init__( metadata: ArrayV2Metadata | ArrayV2MetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> None: ... @overload @@ -416,6 +425,7 @@ def __init__( metadata: ArrayV3Metadata | ArrayMetadataJSON_V3, store_path: StorePath, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> None: ... def __init__( @@ -423,8 +433,9 @@ def __init__( 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)) # An AsyncArray is the resolving layer: build its codec pipeline now, so # pipeline-construction warnings fire at creation, as they always have. _ = self._io._resolved_codec_pipeline @@ -4607,13 +4618,17 @@ 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: @@ -4621,7 +4636,7 @@ def build_array_metadata( 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( @@ -4697,6 +4712,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" @@ -5320,6 +5336,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. @@ -5336,7 +5353,9 @@ 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) @@ -5344,7 +5363,7 @@ def _parse_chunk_encoding_v3( # 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, ...] = () @@ -5357,7 +5376,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 @@ -6254,10 +6275,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. It is consumed where + the document resolves into classes — building the data path — so the + document operations never touch it: a name outside the context fails only + when chunks are read or written, even when the default registries know it. + Without one, the default registries are used. """ _store_path: StorePath _config: ArrayConfig + _context: Context | None _source: ArrayDocumentModel | ArrayMetadata _document_cache: ArrayDocumentModel | None _resolved_cache: ArrayMetadata | None @@ -6270,9 +6299,11 @@ def __init__( document: ArrayDocumentModel | ArrayMetadata | ArrayMetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> None: self._store_path = store_path self._config = parse_array_config(config) + self._context = context self._document_cache = None self._resolved_cache = None self._chunk_grid_cache = None @@ -6284,7 +6315,7 @@ def __init__( # A resolved metadata object passes through; a raw document dict is # read eagerly, as repaired when it needs that, warning as `AsyncArray` # always has. - self._source = parse_array_metadata(document, str(store_path)) + self._source = parse_array_metadata(document, str(store_path), context=context) @classmethod async def build( @@ -6306,12 +6337,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, @@ -6327,9 +6361,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( @@ -6339,6 +6374,7 @@ async def open( name: str | None = None, zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> AsyncArrayIO: """Read the metadata document of an existing array and bind to it. @@ -6354,7 +6390,7 @@ async def open( """ store_path = await make_store_path(store, path=name) data = await get_array_metadata(store_path, zarr_format=zarr_format) - return cls(_array_document_model(data), store_path, config) + return cls(_array_document_model(data), store_path, config, context=context) @property def document(self) -> ArrayDocumentModel: @@ -6382,7 +6418,9 @@ def metadata(self) -> ArrayMetadata: data = _document_json(source) if isinstance(source, ZarrV2ArrayMetadata): data["attributes"] = dict(source.attributes) - self._resolved_cache = parse_array_metadata(data, str(self._store_path)) + self._resolved_cache = parse_array_metadata( + data, str(self._store_path), context=self._context + ) return self._resolved_cache @property @@ -6495,7 +6533,7 @@ def _shard_grid_shape(self) -> tuple[int, ...]: def _resolved_codec_pipeline(self) -> CodecPipeline: if self._codec_pipeline_cache is None: self._codec_pipeline_cache = create_codec_pipeline( - metadata=self.metadata, store=self._store_path.store + metadata=self.metadata, store=self._store_path.store, context=self._context ) return self._codec_pipeline_cache @@ -6714,8 +6752,9 @@ def __init__( document: ArrayDocumentModel | ArrayMetadata | ArrayMetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> None: - self._async_io = AsyncArrayIO(document, store_path, config) + self._async_io = AsyncArrayIO(document, store_path, config, context=context) @classmethod def _wrap(cls, async_io: AsyncArrayIO) -> ArrayIO: @@ -6743,6 +6782,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. @@ -6767,6 +6807,7 @@ def build( chunk_key_encoding=chunk_key_encoding, dimension_names=dimension_names, config=config, + context=context, ) ) ) @@ -6779,13 +6820,18 @@ def open( name: str | None = None, zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> ArrayIO: """Read the metadata document 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 diff --git a/src/zarr/core/context.py b/src/zarr/core/context.py index 494945d9f0..e07f0a75f6 100644 --- a/src/zarr/core/context.py +++ b/src/zarr/core/context.py @@ -7,7 +7,8 @@ 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 experimental explicit-IO array layer (`zarr.experimental.ArrayIO` and +`AsyncArrayIO`) takes a context publicly; everywhere else, both remain internal. """ from __future__ import annotations @@ -15,7 +16,12 @@ from dataclasses import dataclass, replace 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 +33,26 @@ 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. + 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 @classmethod def default(cls) -> Context: @@ -48,6 +67,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, ...] diff --git a/src/zarr/core/dtype/__init__.py b/src/zarr/core/dtype/__init__.py index d1dbd6e2c8..91ac7d9be4 100644 --- a/src/zarr/core/dtype/__init__.py +++ b/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,7 +284,7 @@ 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 @@ -282,7 +294,7 @@ def parse_dtype( return VariableLengthUTF8() # 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: diff --git a/src/zarr/core/metadata/v2.py b/src/zarr/core/metadata/v2.py index 136267fc58..6be1455ee2 100644 --- a/src/zarr/core/metadata/v2.py +++ b/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") diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 2d5d169e6b..74ebbe8c35 100644 --- a/src/zarr/core/metadata/v3.py +++ b/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] diff --git a/src/zarr/experimental/__init__.py b/src/zarr/experimental/__init__.py index f7caaf96a1..097d142c54 100644 --- a/src/zarr/experimental/__init__.py +++ b/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"] diff --git a/src/zarr/registry.py b/src/zarr/registry.py index 307c323153..332f645cfc 100644 --- a/src/zarr/registry.py +++ b/src/zarr/registry.py @@ -25,6 +25,7 @@ from zarr.core.buffer import Buffer, NDBuffer from zarr.core.chunk_key_encodings import ChunkKeyEncoding from zarr.core.common import JSON, ZarrFormat + from zarr.core.context import Context __all__ = [ "Registry", @@ -387,15 +388,22 @@ def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]: return selected_codec_cls -def _resolve_codec(data: dict[str, JSON]) -> Codec: +def _resolve_codec(data: dict[str, JSON], *, context: Context | None = None) -> Codec: """ Get a codec instance from a dict representation of that codec. + + With a `context`, the codec name resolves to a class through it instead of the + default codec registry. """ + name = data["name"] + codec_cls = get_codec_class(name) if context is None else context.codec_class(name) # type: ignore[arg-type] # TODO: narrow the type of the input to only those dicts that map on to codec class instances. - return get_codec_class(data["name"]).from_dict(data) # type: ignore[arg-type] + return codec_cls.from_dict(data) -def _parse_bytes_bytes_codec(data: dict[str, JSON] | Codec) -> BytesBytesCodec: +def _parse_bytes_bytes_codec( + data: dict[str, JSON] | Codec, *, context: Context | None = None +) -> BytesBytesCodec: """ Normalize the input to a `BytesBytesCodec` instance. If the input is already a `BytesBytesCodec`, it is returned as is. If the input is a dict, it @@ -404,7 +412,7 @@ def _parse_bytes_bytes_codec(data: dict[str, JSON] | Codec) -> BytesBytesCodec: from zarr.abc.codec import BytesBytesCodec if isinstance(data, dict): - result = _resolve_codec(data) + result = _resolve_codec(data, context=context) if not isinstance(result, BytesBytesCodec): msg = f"Expected a dict representation of a BytesBytesCodec; got a dict representation of a {type(result)} instead." raise TypeError(msg) @@ -415,7 +423,9 @@ def _parse_bytes_bytes_codec(data: dict[str, JSON] | Codec) -> BytesBytesCodec: return result -def _parse_array_bytes_codec(data: dict[str, JSON] | Codec) -> ArrayBytesCodec: +def _parse_array_bytes_codec( + data: dict[str, JSON] | Codec, *, context: Context | None = None +) -> ArrayBytesCodec: """ Normalize the input to a `ArrayBytesCodec` instance. If the input is already a `ArrayBytesCodec`, it is returned as is. If the input is a dict, it @@ -424,7 +434,7 @@ def _parse_array_bytes_codec(data: dict[str, JSON] | Codec) -> ArrayBytesCodec: from zarr.abc.codec import ArrayBytesCodec if isinstance(data, dict): - result = _resolve_codec(data) + result = _resolve_codec(data, context=context) if not isinstance(result, ArrayBytesCodec): msg = f"Expected a dict representation of an ArrayBytesCodec; got a dict representation of a {type(result)} instead." raise TypeError(msg) @@ -435,7 +445,9 @@ def _parse_array_bytes_codec(data: dict[str, JSON] | Codec) -> ArrayBytesCodec: return result -def _parse_array_array_codec(data: dict[str, JSON] | Codec) -> ArrayArrayCodec: +def _parse_array_array_codec( + data: dict[str, JSON] | Codec, *, context: Context | None = None +) -> ArrayArrayCodec: """ Normalize the input to a `ArrayArrayCodec` instance. If the input is already a `ArrayArrayCodec`, it is returned as is. If the input is a dict, it @@ -444,7 +456,7 @@ def _parse_array_array_codec(data: dict[str, JSON] | Codec) -> ArrayArrayCodec: from zarr.abc.codec import ArrayArrayCodec if isinstance(data, dict): - result = _resolve_codec(data) + result = _resolve_codec(data, context=context) if not isinstance(result, ArrayArrayCodec): msg = f"Expected a dict representation of an ArrayArrayCodec; got a dict representation of a {type(result)} instead." raise TypeError(msg) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 68f60c2e5f..b32bf6e104 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -10,12 +10,24 @@ from zarr import ArrayIO, AsyncArrayIO from zarr.api.asynchronous import open_array +from zarr.codecs import BytesCodec, ZstdCodec from zarr.core.buffer import default_buffer_prototype +from zarr.core.codec_pipeline import BatchedCodecPipeline +from zarr.core.dtype import UInt16, data_type_registry +from zarr.core.dtype.registry import DataTypeRegistry from zarr.errors import ArrayNotFoundError, UnknownCodecError +from zarr.experimental import Context from zarr.storage import MemoryStore if TYPE_CHECKING: - from zarr.core.common import ZarrFormat + from zarr.core.common import JSON, ZarrFormat + +ZSTD_JSON: dict[str, JSON] = {"name": "zstd", "configuration": {"level": 0, "checksum": False}} + + +class _MarkedPipeline(BatchedCodecPipeline): + """A pipeline class distinguishable from the default, for context tests.""" + METADATA_DOCUMENTS = ("zarr.json", ".zarray", ".zattrs", ".zgroup") @@ -202,3 +214,64 @@ async def test_erase_chunk_out_of_grid_raises() -> None: await aio.erase_chunk((2, 0)) with pytest.raises(IndexError, match="outside the chunk grid"): await aio.erase_chunk((0,)) + + +async def test_context_scopes_resolution() -> None: + """A context defines the data types, codecs and pipeline in scope: names + resolve through it when the data path is built, and the pipeline class is + honored.""" + store = MemoryStore() + data = np.arange(16, dtype="uint16").reshape(4, 4) + registry = DataTypeRegistry() + registry.register(UInt16._zarr_v3_name, UInt16) + context = Context( + data_types=registry, + codecs={"bytes": BytesCodec, "zstd": ZstdCodec}, + codec_pipeline=_MarkedPipeline, + ) + aio = await AsyncArrayIO.build( + store, + name="foo", + shape=data.shape, + dtype="uint16", + chunks=(2, 2), + compressors=ZSTD_JSON, + context=context, + ) + assert isinstance(aio.to_array().codec_pipeline, _MarkedPipeline) + await aio.setitem(Ellipsis, data) + await aio.store_metadata() + + opened = await AsyncArrayIO.open(store, name="foo", context=context) + np.testing.assert_array_equal(await opened.getitem(Ellipsis), data) + + +async def test_context_unknown_codec_resolution_raises() -> None: + """A stored codec name outside the context opens for document work and + fails to resolve when the data path is built, naming the scope.""" + store = MemoryStore() + aio = await AsyncArrayIO.build(store, name="foo", shape=(4,), dtype="uint8") + await aio.store_metadata() + context = Context(data_types=data_type_registry, codecs={"bytes": BytesCodec}) + opened = await AsyncArrayIO.open(store, name="foo", context=context) + await opened.verify() # document work needs no resolution + with pytest.raises(UnknownCodecError, match="Codecs in scope: \\['bytes'\\]"): + await opened.getitem(Ellipsis) + + +async def test_context_unknown_codec_on_build_raises() -> None: + context = Context(data_types=data_type_registry, codecs={"bytes": BytesCodec}) + with pytest.raises(UnknownCodecError, match="'zstd' is not in this context"): + await AsyncArrayIO.build( + MemoryStore(), + shape=(4,), + dtype="uint8", + compressors=ZSTD_JSON, + context=context, + ) + + +async def test_context_unknown_dtype_raises() -> None: + context = Context(data_types=DataTypeRegistry()) + with pytest.raises(ValueError, match="No Zarr data type found"): + await AsyncArrayIO.build(MemoryStore(), shape=(4,), dtype="uint16", context=context)