From 8c3d1eedcb52c3b6fa3019a361e2c9284e1dd70d Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 16:02:23 +0100 Subject: [PATCH] feat(context): scope ArrayIO's extension resolution with a Context Extends zarr.core.context.Context (exported as zarr.experimental.Context) with `codecs` and `codec_pipeline`, makes `data_types` optional, and threads an optional context through data type and codec resolution, both formats' metadata `from_dict`, `parse_array_metadata`, the per-format builders, the repair check before chunk writes, and ArrayIO/AsyncArrayIO/AsyncArray (with `with_context`). A kind the context leaves unset resolves from the defaults; a kind it sets is a closed world over top-level names (`limits_names`), checking instances, structured dtypes' fields and default codec chains as well as names, and re-resolving only the kinds it limits. A context that only selects the pipeline, or sets nothing, resolves as no context does. The context's pipeline choice fills the slot the global configuration fills otherwise, so the context is the per-binding replacement for that global state. Bindings keep their context through derivation, AsyncArray's pickled state, with_config and dataclasses.replace, and compare it in equality; the pipeline-parity, indexing and fork-pool tests select their pipeline through a context instead of global configuration. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- changes/420.feature.md | 4 + src/zarr/core/array.py | 548 ++++++++-- src/zarr/core/context.py | 173 ++- src/zarr/core/dtype/__init__.py | 116 +- src/zarr/core/dtype/common.py | 11 + src/zarr/core/dtype/registry.py | 2 +- src/zarr/core/metadata/io.py | 18 +- src/zarr/core/metadata/v2.py | 17 +- src/zarr/core/metadata/v3.py | 130 ++- src/zarr/experimental/__init__.py | 3 +- tests/test_array.py | 33 +- tests/test_array_io.py | 1700 ++++++++++++++++++++++++++++- tests/test_indexing.py | 16 +- tests/test_pipeline_parity.py | 113 +- 14 files changed, 2646 insertions(+), 238 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..7a61840155 --- /dev/null +++ b/changes/420.feature.md @@ -0,0 +1,4 @@ +Added a `context` (`zarr.experimental.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 (for codecs, with an error naming the codecs in scope). The codec scope governs Zarr format 3 codecs; a Zarr format 2 array's filters and compressor are numcodecs objects, which no context governs. +`AsyncArray.context` gives an array's context (`Array.io.context` for a synchronous array), and `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; a binding opened from the store resolves no names until its data path is built, so it fails there. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 867a2bdd6b..8ad994cb08 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -21,6 +21,7 @@ ClassVar, Final, Literal, + Protocol, Required, TypedDict, Unpack, @@ -87,12 +88,17 @@ product, ) from zarr.core.config import config as zarr_config +from zarr.core.context import Context from zarr.core.dtype import ( Structured, VariableLengthBytes, VariableLengthUTF8, ZDType, ZDTypeLike, + _json_to_reread, + _matched_in_scope, + _with_byte_order_of, + get_data_type_from_json, parse_dtype, ) from zarr.core.dtype.common import HasEndianness, HasItemSize, HasObjectCodec @@ -156,7 +162,10 @@ RectilinearChunkGridMetadataJSON, RegularChunkGridMetadata, RegularChunkGridMetadataJSON, + _codec_in_scope, + _rescoped, create_chunk_grid_metadata, + parse_codecs, parse_node_type_array, ) from zarr.core.sync import sync @@ -242,26 +251,33 @@ 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`, stored + extension names resolve through it instead of the default registries; a metadata + object is already resolved and passes through any context (an `AsyncArrayIO` + binding re-resolves one it is given under its 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 are (see `zarr.core.metadata.repair`), so an array can be built from it. No data was read or written under that chunk size, so the reading is silent.""" if isinstance(data, ArrayV2Metadata) and 0 in data.chunks: - return parse_stored_array(data.to_buffer_dict(default_buffer_prototype()), 2) + return parse_stored_array( + data.to_buffer_dict(default_buffer_prototype()), 2, context=context + ) if isinstance(data, ArrayMetadata): return data 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) _reject_storage_transformers(meta_out) 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( @@ -453,6 +469,8 @@ def __init__( metadata: ArrayV2Metadata | ArrayV2MetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, + *, + context: Context | None = None, ) -> None: ... @overload @@ -461,6 +479,8 @@ def __init__( metadata: ArrayV3Metadata | ArrayMetadataJSON_V3, store_path: StorePath, config: ArrayConfigLike | None = None, + *, + context: Context | None = None, ) -> None: ... def __init__( @@ -468,16 +488,30 @@ def __init__( metadata: ArrayMetadata | ArrayMetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, + *, + context: Context | None = None, ) -> None: - # An AsyncArray reads metadata into the legacy classes, as it always has. - metadata_parsed = parse_array_metadata(metadata, str(store_path)) if isinstance(metadata, dict): - # A document given as a dict stands for what the store holds, as the - # document an array is opened from does. + # An AsyncArray reads a dict into the legacy classes, as it always has, + # once, under its context; the dict stands for what the store holds, as + # the document an array is opened from does. + metadata_parsed = parse_array_metadata(metadata, str(store_path), context=context) _record_given_document(store_path, metadata, metadata_parsed.zarr_format) - object.__setattr__( - self, "_io", AsyncArrayIO._from_resolved(metadata_parsed, store_path, config) - ) + io = AsyncArrayIO._from_resolved(metadata_parsed, store_path, config, context) + elif ( + context is not None + and context.limits_names + and isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata) + ): + # Re-resolved under the context, so the scope checks its names; the + # object itself is bound when the scope changes nothing, as with no + # context. + resolved = _legacy_resolved_in_context(metadata, str(store_path), context) + io = AsyncArrayIO._from_resolved(resolved, store_path, config, context) + else: + # The object itself is bound, as AsyncArray always bound it. + io = AsyncArrayIO._from_resolved(metadata, store_path, config, context) + object.__setattr__(self, "_io", io) # 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 @@ -496,9 +530,22 @@ def io(self) -> AsyncArrayIO: append gives the array a binding of its own, so a binding taken before keeps the shape it had; an attribute update changes the metadata in place, as it always has, so it shows through a binding taken before, and through copies - of the array that share its metadata.""" + of the array that share its metadata. + + An array made by the legacy API (`zarr.open_array`, `zarr.create_array`) or + by `dataclasses.replace` holds legacy metadata, so this binding's `metadata` + is that metadata as zarr writes it, which may differ from the stored document + (a NaN payload, a codec's configuration): `ArrayIO.open(...).to_array()` + keeps the document as stored.""" return self._io + @property + def context(self) -> Context | None: + """The context this array resolves in, or `None` for the defaults. It is not a + dataclass field: `dataclasses.replace` builds from the fields, so its result + has no context; `with_config` keeps it.""" + return self._io._context + def __getstate__(self) -> dict[str, object] | tuple[dict[str, object], dict[str, object]]: # The pickled state is the fields', as it has always been, so a pickle # loads in either version; the binding's own state rides along under one @@ -520,7 +567,10 @@ def __getstate__(self) -> dict[str, object] | tuple[dict[str, object], dict[str, "_chunk_grid": self._chunk_grid, "config": self.config, } - binding_state: dict[str, object] = {"pipeline_class": self._io._pipeline_class} + binding_state: dict[str, object] = { + "pipeline_class": self._io._pipeline_class, + "context": self._io._context, + } if type(self._io) is not AsyncArrayIO: # Carried only when it is a subclass, so the pickle stays loadable where # that class is not importable only if it is used. @@ -567,6 +617,7 @@ def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, A metadata if model is None else model, state["store_path"], state["config"], + binding_state.get("context"), pipeline_class, ) if model is not None: @@ -1364,6 +1415,12 @@ def with_config(self, config: ArrayConfigLike) -> Self: Returns ------- A new Array + + Notes + ----- + A subclass's own `__init__` runs without a context, as it always has, so + under a context it builds the globally configured pipeline first: that + pipeline must be buildable, though the array then runs its context's. """ if isinstance(config, ArrayConfig): new_config = config @@ -1371,11 +1428,29 @@ def with_config(self, config: ArrayConfigLike) -> Self: # Merge new config with existing config, so missing keys are inherited # from the current array rather than from global defaults new_config = ArrayConfig(**{**self.config.to_dict(), **config}) # type: ignore[arg-type] - # Built through the constructor, as it always was, so a subclass's own - # __init__ runs; then over this array's kind of binding (its class, and - # its model, so the document stays as stored). - rebuilt = type(self)(metadata=self.metadata, store_path=self.store_path, config=new_config) - object.__setattr__(rebuilt, "_io", self._io._rebound_like(rebuilt._io, rebuilt.metadata)) + # Then bound over this array's kind of binding: its class, context and + # model, so the document stays as stored. + binding = self._io + if binding._context is not None and type(self).__init__ is AsyncArray.__init__: + # Under a context, AsyncArray's own __init__ would bind this metadata + # object as it is: bind it so directly, checking no names again and + # building no other pipeline. Without one, the constructor runs, as it + # always has. + rebuilt = object.__new__(type(self)) + built = AsyncArrayIO._from_resolved( + self.metadata, self.store_path, new_config, binding._context + ) + kept = self.metadata + else: + # The constructor runs, as it always has; a subclass's own __init__ runs + # without a context (one written before contexts takes none), so under + # one it builds the globally configured pipeline before the array is + # bound under its context. + rebuilt = type(self)( + metadata=self.metadata, store_path=self.store_path, config=new_config + ) + built, kept = rebuilt._io, rebuilt.metadata + object.__setattr__(rebuilt, "_io", binding._rebound_like(built, kept)) return rebuilt async def nchunks_initialized(self) -> int: @@ -4950,11 +5025,12 @@ def _array_metadata_v2( filters: Iterable[dict[str, JSON] | Numcodec] | None = None, compressor: CompressorLikev2 = None, attributes: dict[str, JSON] | None = None, + context: Context | None = None, ) -> ArrayV2Metadata: """`build_array_metadata_v2` without its object codec check, which the callers that create arrays from user arguments make (or, for the legacy `create` path, have never made) themselves.""" - zdtype = parse_dtype(dtype, zarr_format=2) + zdtype = parse_dtype(dtype, zarr_format=2, context=context) if order is None: order = zarr_config.get("array.order") if fill_value is DEFAULT_FILL_VALUE: @@ -4987,8 +5063,10 @@ def build_array_metadata_v2( filters: Iterable[dict[str, JSON] | Numcodec] | None = None, compressor: CompressorLikev2 = None, attributes: dict[str, JSON] | None = None, + context: Context | None = None, ) -> ArrayV2Metadata: - """Build Zarr format 2 array metadata from its fields. + """Build Zarr format 2 array metadata from its fields. With a `context`, + the data type resolves through it instead of the default registry. Each keyword is a `.zarray` field; `attributes` is the `.zattrs` one. `fill_value` defaults to the data type's default scalar (an explicit @@ -5006,6 +5084,7 @@ def build_array_metadata_v2( filters=filters, compressor=compressor, attributes=attributes, + context=context, ) _check_object_codec_v2(metadata.dtype, metadata.filters, metadata.compressor) return metadata @@ -5021,8 +5100,11 @@ def build_array_metadata_v3( codecs: Iterable[Codec | dict[str, JSON]] | None = None, dimension_names: DimensionNamesLike = None, attributes: dict[str, JSON] | None = None, + context: Context | None = None, ) -> ArrayV3Metadata: - """Build Zarr format 3 array metadata from its fields. + """Build Zarr format 3 array metadata from its fields. With a `context`, + the data type and top-level codec names resolve through it instead of the + default registries. Each keyword is a `zarr.json` field. `chunk_grid` takes grid metadata, its JSON, or a chunk shape, which declares a regular grid. `codecs` @@ -5030,7 +5112,7 @@ def build_array_metadata_v3( the default encoding with `/`; `fill_value` defaults to the data type's default scalar (`None` means the default too: the field is required). """ - zdtype = parse_dtype(data_type, zarr_format=3) + zdtype = parse_dtype(data_type, zarr_format=3, context=context) grid: ChunkGridMetadata | dict[str, JSON] if isinstance(chunk_grid, RegularChunkGridMetadata | RectilinearChunkGridMetadata): grid = chunk_grid @@ -5041,13 +5123,26 @@ def build_array_metadata_v3( filters: tuple[ArrayArrayCodec, ...] compressors: tuple[BytesBytesCodec, ...] + codecs_parsed: tuple[Codec | dict[str, JSON], ...] if codecs is None: # no data types have default filters filters = () serializer = default_serializer_v3(zdtype) compressors = default_compressors_v3(zdtype) - codecs_parsed = (*filters, serializer, *compressors) + default_chain: tuple[Codec, ...] = (*filters, serializer, *compressors) + codecs_parsed = default_chain + if context is not None and context.codecs is not None: + # The defaults resolve under the context too: a closed codec scope + # must not store a chain it could not read back. + codecs_parsed = parse_codecs( + tuple(_codec_in_scope(codec, context) for codec in default_chain), + context=context, + ) + elif context is not None and context.codecs is not None: + # Codec names resolve here, under the context; the metadata's own + # `parse_codecs` passes the resolved instances through untouched. + codecs_parsed = parse_codecs(tuple(codecs), context=context) else: codecs_parsed = tuple(codecs) @@ -5064,16 +5159,17 @@ def build_array_metadata_v3( else: fill_value_parsed = fill_value - return ArrayV3Metadata( + metadata = ArrayV3Metadata( shape=parse_shapelike(shape), data_type=zdtype, chunk_grid=grid, chunk_key_encoding=chunk_key_encoding_parsed, fill_value=fill_value_parsed, - codecs=codecs_parsed, # type: ignore[arg-type] + codecs=codecs_parsed, dimension_names=tuple(dimension_names) if dimension_names is not None else None, attributes=attributes if attributes is not None else {}, ) + return _rescoped(metadata, context) async def create_array( @@ -5298,7 +5394,39 @@ def _sharding_codec(metadata: ArrayMetadata) -> ShardingCodec | None: return None -async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadata) -> None: +_SCOPED_READ_ERRORS: Final = (ValueError, TypeError, KeyError) +"""What a document read under a scope raises when the scope cannot read it: a name +outside it, or a scope class that refuses the configuration the document holds. A +codec's `from_dict` refuses an unexpected keyword with `TypeError` and a missing one +with `KeyError` as readily as a bad value with `ValueError`, so all three are a +scope that cannot read the document; the defaults then read it, as with no context.""" + + +class _ScopedRead[T](Protocol): + """A read of stored metadata under a context.""" + + def __call__(self, *, context: Context | None) -> T: ... + + +def _read_in_scope_or_defaults[T](read: _ScopedRead[T], context: Context | None) -> T: + """`read(context=...)` under `context`. A document a name-limiting scope cannot + read is read under the defaults, as with no context: its layout decides, not its + names. One they cannot read either raises the scope's error, which names what it + lacks.""" + try: + return read(context=context) + except _SCOPED_READ_ERRORS as scoped: + if context is None or not context.limits_names: + raise + try: + return read(context=None) + except _SCOPED_READ_ERRORS: + raise scoped from None + + +async def _store_repaired_document( + store_path: StorePath, metadata: ArrayMetadata, context: Context | None = None +) -> None: """Store the repair of the array document at `store_path`, if it needs one that moves chunks, before chunks are written under `metadata`. @@ -5341,16 +5469,22 @@ async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadat stored = cache.needing_repair(path, zarr_format) else: stored = cache.copy_needing_repair(path, zarr_format) - if stored is None or _chunk_layout(read_repaired(stored, zarr_format)) != _chunk_layout( - metadata - ): + if stored is None: + return + recorded = _read_in_scope_or_defaults( + functools.partial(read_repaired, stored, zarr_format), context + ) + if _chunk_layout(recorded) != _chunk_layout(metadata): return documents = await read_documents(store_path, keys) document = stored_array_document(documents, zarr_format) if document is None: cache.observe(path, None, zarr_format) return - current = parse_stored_array(documents, zarr_format, str(store_path)) + current = _read_in_scope_or_defaults( + functools.partial(parse_stored_array, documents, zarr_format, str(store_path)), + context, + ) if _chunk_layout(current) != _chunk_layout(metadata): raise ValueError( f"Array {str(store_path)!r}: the metadata stored has changed since " @@ -6546,6 +6680,9 @@ async def _delete_key(store_path: StorePath, key: str) -> None: await (store_path / key).delete() +_NO_CONTEXT: Final = Context() +"""A context that sets nothing, which binds as no context.""" + _CARRIED_KEYS: Final = ("_stored_entry", "_stored_copy") """The keys a pickled binding carries its store cache's entry under.""" @@ -6565,7 +6702,7 @@ async def _delete_key(store_path: StorePath, key: str) -> None: def _bound_metadata_dict( - data: dict[str, Any], store_path: StorePath + data: dict[str, Any], store_path: StorePath, context: Context | None ) -> ArrayMetadataModel | ArrayMetadata: """A metadata dict given to a binding: read as a model when it is valid JSON metadata, so names resolve only when chunks are read or written; @@ -6581,7 +6718,7 @@ def _bound_metadata_dict( zarr_format: ZarrFormat = 2 if data.get("zarr_format") == 2 else 3 if data.get("storage_transformers"): # Bound as legacy metadata, which rejects them, as AsyncArray always has. - return parse_array_metadata(data, path) + return parse_array_metadata(data, path, context=context) if not moves_chunks(data, zarr_format): # A document zarr reads as repaired, moving chunks, is never bound as a # model, also where the model accepts it. @@ -6589,7 +6726,7 @@ def _bound_metadata_dict( return _array_metadata_model(data) except MetadataValidationError: pass - legacy = parse_array_metadata(data, path) + legacy = parse_array_metadata(data, path, context=context) # An invalid document given as a dict stands for what the store holds, as # one given to AsyncArray does. _record_given_document(store_path, data, legacy.zarr_format) @@ -6678,6 +6815,85 @@ def _with_legacy_mark(metadata: ArrayMetadata, mark: ArrayDocument | None) -> Ar return marked +def _legacy_resolved_in_context( + source: ArrayMetadata, path: str, context: Context | None +) -> ArrayMetadata: + """`source` as `parse_array_metadata` takes a metadata object: v2 metadata + with a chunk size of 0 is read as repaired, silently, under `context`, and + any other object is returned itself, or, when `context` limits names, with + each kind it limits resolved through it (see `_checked_in_scope`), so the scope + checks those names.""" + v2_chunk_zero = isinstance(source, ArrayV2Metadata) and 0 in source.chunks + if context is None or not context.limits_names or v2_chunk_zero: + return parse_array_metadata(source, path, context=context) + return _checked_in_scope(source, context) + + +def _dtype_field(source: ArrayMetadata) -> str: + return "dtype" if isinstance(source, ArrayV2Metadata) else "data_type" + + +def _checked_in_scope(source: ArrayMetadata, context: Context) -> ArrayMetadata: + """`source` with each kind `context` limits resolved through it, so an instance + outside the scope fails and the scope's class for a name is the one bound; the + kinds it leaves unset are kept. A copy (with its attributes copied) when that + changes a class, else `source` itself.""" + zarr_format = source.zarr_format + changes: dict[str, Any] = {} + if context.data_types is not None: + dtype = _matched_in_scope(source.dtype, context, zarr_format) + if dtype != source.dtype: + changes[_dtype_field(source)] = dtype + if context.codecs is not None and isinstance(source, ArrayV3Metadata): + # Each is the scope's class for the name it writes once `_codec_in_scope` + # returns it. + codecs = tuple(_codec_in_scope(codec, context) for codec in source.codecs) + if any(new is not old for new, old in zip(codecs, source.codecs, strict=True)): + changes["codecs"] = codecs + if not changes: + # Everything is the scope's already: the object itself, shared as one + # given with no context is. + return source + return replace(source, attributes=dict(source.attributes), **changes) + + +def _reread_in_defaults(source: ArrayMetadata, *, data_types: bool, codecs: bool) -> ArrayMetadata: + """A copy of `source` with its data type (`data_types`) and its top-level codecs + (`codecs`, Zarr format 3 only) read again from their JSON in the default + registries, dropping the classes a scope bound; the rest, and the attributes + (copied), are kept.""" + zarr_format = source.zarr_format + changes: dict[str, Any] = {"attributes": dict(source.attributes)} + if data_types: + reread = get_data_type_from_json( + _json_to_reread(source.dtype, zarr_format), zarr_format=zarr_format + ) + changes[_dtype_field(source)] = _with_byte_order_of(reread, source.dtype) + if codecs and isinstance(source, ArrayV3Metadata): + # An instance of the default class for its name is kept whole (with the + # codecs nested in it, which no scope governed). + changes["codecs"] = parse_codecs( + [_codec_in_scope(codec, _NO_CONTEXT) for codec in source.codecs] + ) + return replace(source, **changes) + + +def _context_kwargs(context: Context | None) -> dict[str, Context]: + """`context` as a keyword for a binding class's constructor, or none for `None`: + a subclass constructor written without one still builds and opens with no + context, and one given (`Context()` included) reaches it as given. A subclass whose + constructor defaults `context` to a scope gets that scope only when the + metadata it is given is read there: a builder declares under the context + passed to it.""" + return {} if context is None else {"context": context} + + +def _pipeline_class_for(context: Context | None) -> type[CodecPipeline]: + """The codec pipeline class a binding made now uses: the context's choice, + or else the one the global configuration names.""" + return get_pipeline_class() if context is None else context.pipeline_class() + + def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: """A copy of a legacy metadata object with its own attributes dict, so an attribute update through one binding (which replaces top-level keys in @@ -6853,10 +7069,13 @@ class AsyncArrayIO: the legacy classes (`metadata_legacy`) on first use. The bound metadata is assumed, never checked against the store, and an array opened from metadata that needed a repair stores that repair before its first chunk write. The codec pipeline implementation is - chosen from the global configuration when the binding is made; the - pipeline itself, with any options it reads from the configuration, is - built at the first data operation. `AsyncArray` wraps an `AsyncArrayIO` - and delegates its metadata and selection IO to it. + chosen when the binding is made, by its `context` or else the global + configuration; the pipeline itself, with any options it reads from the + configuration, is built at the first data operation. Under a `context`, a + resolved legacy metadata object given to the constructor or to + `with_metadata` is re-resolved there, so the scope checks its names. + `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and + selection IO to it. Every way a binding is made or changed (construction, the builders, `open`, the `with_*` methods, pickling, and the resize or attribute update @@ -6872,6 +7091,7 @@ class AsyncArrayIO: _store_path: StorePath _config: ArrayConfig _pipeline_class: type[CodecPipeline] + _context: Context | None _source: ArrayMetadataModel | ArrayMetadata _resolved_cache: ArrayMetadata | None _chunk_grid_cache: ChunkGrid | None @@ -6884,37 +7104,47 @@ def __init__( metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, store_path: StorePath, config: ArrayConfigLike | None = None, + *, + context: Context | None = None, ) -> None: source: ArrayMetadataModel | ArrayMetadata - if isinstance(metadata, ArrayMetadata): - # A resolved object is read (v2 with a chunk size of 0 as repaired) - # and never shared with the caller. - source = parse_array_metadata(metadata, str(store_path)) + if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): + # A resolved object given from outside is re-resolved under a context + # that limits names, so the scope checks them; otherwise it is read + # (v2 with a chunk size of 0 as repaired). Either way it is never + # shared with the caller. + source = _legacy_resolved_in_context(metadata, str(store_path), context) if source is metadata: source = _copied_legacy(source) elif isinstance(metadata, _model_classes()): source = _checked_model(metadata) elif isinstance(metadata, Mapping): - source = _bound_metadata_dict(cast("dict[str, Any]", metadata), store_path) + source = _bound_metadata_dict(cast("dict[str, Any]", metadata), store_path, context) else: raise TypeError( "Expected array metadata (a zarr_metadata model, ArrayV2Metadata or " f"ArrayV3Metadata) or a metadata mapping, got {type(metadata).__name__}." ) - # The pipeline implementation is chosen when the binding is made, from - # the global configuration, exactly as `AsyncArray` has always chosen - # it; the pipeline itself is built on first use. - self._bind(source, store_path, parse_array_config(config), get_pipeline_class()) + # Without a context, the pipeline implementation is chosen when the + # binding is made, from the global configuration, exactly as + # `AsyncArray` has always chosen it; a context's choice wins. The + # pipeline itself is built on first use. + self._bind( + source, store_path, parse_array_config(config), context, _pipeline_class_for(context) + ) def _bind( self, source: ArrayMetadataModel | ArrayMetadata, store_path: StorePath, config: ArrayConfig, + context: Context | None, pipeline_class: type[CodecPipeline], ) -> None: self._store_path = store_path self._config = config + # A context that sets nothing is no context: they resolve alike. + self._context = None if context == _NO_CONTEXT else context self._pipeline_class = pipeline_class self._source = source self._resolved_cache = None @@ -6925,17 +7155,20 @@ def _bind( @classmethod def _from_resolved( - cls, source: ArrayMetadata, store_path: StorePath, config: ArrayConfigLike | None + cls, + source: ArrayMetadata, + store_path: StorePath, + config: ArrayConfigLike | None, + context: Context | None, ) -> Self: - """A binding over `source` itself, already resolved (as `AsyncArray` binds the - metadata it is given): read as `parse_array_metadata` reads a metadata - object, but not copied.""" + """A binding over `source`, which was just resolved under `context`, so its + names are not resolved again; it is still read as `parse_array_metadata` + reads a metadata object (v2 metadata with a chunk size of 0 as repaired, + under `context`).""" + source = parse_array_metadata(source, str(store_path), context=context) binding = object.__new__(cls) binding._bind( - parse_array_metadata(source, str(store_path)), - store_path, - parse_array_config(config), - get_pipeline_class(), + source, store_path, parse_array_config(config), context, _pipeline_class_for(context) ) return binding @@ -6946,6 +7179,7 @@ def build_v2( *, name: str | None = None, config: ArrayConfigLike | None = None, + context: Context | None = None, **kwargs: Unpack[ArrayV2DeclarationKwargs], ) -> Self: """Declare a Zarr format 2 array bound to `store`, storing nothing. @@ -6954,7 +7188,12 @@ def build_v2( `store` must be a `Store` or `StorePath`: coercing strings opens stores. `store_metadata` stores the result. """ - return cls(build_array_metadata_v2(**kwargs), _bound_store_path(store, name), config) + return cls( + build_array_metadata_v2(context=context, **kwargs), + _bound_store_path(store, name), + config, + **_context_kwargs(context), + ) @classmethod def build_v3( @@ -6963,6 +7202,7 @@ def build_v3( *, name: str | None = None, config: ArrayConfigLike | None = None, + context: Context | None = None, **kwargs: Unpack[ArrayV3DeclarationKwargs], ) -> Self: """Declare a Zarr format 3 array bound to `store`, storing nothing. @@ -6971,7 +7211,12 @@ def build_v3( `build_array_metadata_v3`). `store` must be a `Store` or `StorePath`: coercing strings opens stores. `store_metadata` stores the result. """ - return cls(build_array_metadata_v3(**kwargs), _bound_store_path(store, name), config) + return cls( + build_array_metadata_v3(context=context, **kwargs), + _bound_store_path(store, name), + config, + **_context_kwargs(context), + ) @classmethod async def open( @@ -6981,6 +7226,7 @@ async def open( name: str | None = None, zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> Self: """Read the stored metadata of an existing array and bind to it. @@ -6993,7 +7239,7 @@ async def open( mode: AccessModeLiteral | None = "r+" if isinstance(store, str | os.PathLike) else None store_path = await make_store_path(store, path=name, mode=mode) data = await get_array_metadata(store_path, zarr_format=zarr_format) - binding = cls(_array_metadata_model(data), store_path, config) + binding = cls(_array_metadata_model(data), store_path, config, **_context_kwargs(context)) observe_read(store_path, data) return binding @@ -7031,7 +7277,7 @@ def metadata_legacy(self) -> ArrayMetadata: data = _metadata_model_json(source) if isinstance(source, _model_classes()[0]): data.setdefault("attributes", {}) - resolved = parse_array_metadata(data, str(self._store_path)) + resolved = parse_array_metadata(data, str(self._store_path), context=self._context) # Threads resolving at once all keep the first object stored: an # `AsyncArray` changes its attributes in place, so they must share it. # (A rebind of this binding racing the read is a write racing a read, @@ -7055,33 +7301,55 @@ def _detached(self) -> Self: def _rebound_like(self, built: AsyncArrayIO, metadata: ArrayMetadata) -> Self: """How an array rebuilt through its constructor (`built`, which bound - `metadata`) keeps the binding it had: a binding of this binding's class, at - `built`'s location, with its configuration, pipeline class and caches, over - this binding's model when `metadata` is still this binding's resolved form - (shared), else over `metadata` itself (a subclass `__init__` substituted it). - A `_bind` that runs another pipeline class wins over any pipeline `built` - holds, also one a subclass `__init__` assigned.""" + `metadata`) keeps the binding it had: a binding of this binding's class and + context, at `built`'s location and configuration, over this binding's model + when `metadata` is still this binding's resolved form (shared), else over + `metadata` itself (a subclass `__init__` substituted it). Its pipeline class + is `built`'s, chosen anew as a rebuilt array always chose it, unless this + binding's context chose it. A `_bind` that runs another pipeline class wins + over any pipeline `built` holds, also one a subclass `__init__` assigned.""" rebound = object.__new__(type(self)) source = self._source - keeps_model = ( - not isinstance(source, ArrayV2Metadata | ArrayV3Metadata) - and metadata is self._resolved_if_any - ) + context = self._context + substituted = metadata is not self._resolved_if_any + keeps_model = not substituted and not isinstance(source, ArrayV2Metadata | ArrayV3Metadata) + if substituted and context is not None and context.limits_names: + # Resolved under this binding's context, as a constructor given it + # under that context resolves it. + metadata = _legacy_resolved_in_context(metadata, str(built._store_path), context) + chosen_anew = context is None or context.codec_pipeline is None + pipeline_class = built._pipeline_class if chosen_anew else self._given_pipeline_class() rebound._bind( source if keeps_model else metadata, built._store_path, built._config, - built._pipeline_class, + self._context, + pipeline_class, ) if keeps_model: rebound._resolved_cache = metadata - if built._pipeline_class is rebound._pipeline_class: - # The pipeline the constructor built is the one this binding runs; a - # `_bind` that chose another pipeline class builds its own. + if ( + built._pipeline_class is rebound._pipeline_class + and built._resolved_if_any is rebound._resolved_if_any + ): + # The pipeline the constructor built is the one this binding runs, + # over the same metadata; a `_bind` that chose another pipeline class, + # or metadata resolved again under the context, builds its own. rebound._codec_pipeline_cache = built._codec_pipeline_cache rebound._chunk_grid_cache = built._chunk_grid_cache + # Built now, as the constructor builds it, so its warnings fire here. + _ = rebound._resolved_codec_pipeline return rebound + def _given_pipeline_class(self) -> type[CodecPipeline]: + """The pipeline class a binding made from this one is given: its context's + choice when it made one, else this binding's own (which a `_bind` that + chose another class records, so derived bindings run that class too).""" + context = self._context + if context is not None and context.codec_pipeline is not None: + return context.codec_pipeline + return self._pipeline_class + @property def store_path(self) -> StorePath: """The bound location in a store.""" @@ -7092,6 +7360,12 @@ def config(self) -> ArrayConfig: """The runtime configuration.""" return self._config + @property + def context(self) -> Context | None: + """The context this binding resolves extension names and its codec pipeline + under, or `None` for the defaults; `with_context` rebinds it.""" + return self._context + @property def _zarr_format(self) -> ZarrFormat: source = self._source @@ -7111,30 +7385,30 @@ def _resolved_if_any(self) -> ArrayMetadata | None: def _rebind_at(self, store_path: StorePath, config: ArrayConfig) -> None: """Bind this binding again, through `_bind`, at `store_path` with `config`, keeping its caches: setting an array's location or configuration never - rebuilt its pipeline.""" - caches = ( - self._resolved_cache, - self._chunk_grid_cache, - self._chunk_key_encoding_cache, - self._codec_pipeline_cache, - ) - self._bind(self._source, store_path, config, self._pipeline_class) - ( - self._resolved_cache, - self._chunk_grid_cache, - self._chunk_key_encoding_cache, - self._codec_pipeline_cache, - ) = caches + rebuilt its pipeline (unless `_bind` chose another pipeline class, which is + then built on first use: the store may already hold the change, so a + pipeline that refuses to build fails the next read or write, not this).""" + caches = (self._resolved_cache, self._chunk_grid_cache, self._chunk_key_encoding_cache) + pipeline, pipeline_class = self._codec_pipeline_cache, self._pipeline_class + self._bind(self._source, store_path, config, self._context, self._given_pipeline_class()) + self._resolved_cache, self._chunk_grid_cache, self._chunk_key_encoding_cache = caches + if self._pipeline_class is pipeline_class: + self._codec_pipeline_cache = pipeline def _rebind_resolved( self, resolved: ArrayMetadata, chunk_grid: ChunkGrid | None = None ) -> None: """Make `resolved` this binding's state, through `_bind`. The codec pipeline is kept: rebinding is how `resize` swaps the shape, which leaves the chunk - encoding as it was.""" - pipeline = self._codec_pipeline_cache - self._bind(resolved, self._store_path, self._config, self._pipeline_class) - self._codec_pipeline_cache = pipeline + encoding as it was (unless `_bind` chose another pipeline class, which is + then built on first use: `resize` has already stored the new shape, so a + pipeline that refuses to build fails the next read or write, not the resize).""" + pipeline, pipeline_class = self._codec_pipeline_cache, self._pipeline_class + self._bind( + resolved, self._store_path, self._config, self._context, self._given_pipeline_class() + ) + if self._pipeline_class is pipeline_class: + self._codec_pipeline_cache = pipeline self._chunk_grid_cache = chunk_grid @property @@ -7277,9 +7551,9 @@ def _derive( store_path: StorePath, config: ArrayConfig, ) -> Self: - """A new binding over `source`, keeping this one's captured pipeline.""" + """A new binding over `source`, keeping this one's context and captured pipeline.""" derived = object.__new__(type(self)) - derived._bind(source, store_path, config, self._pipeline_class) + derived._bind(source, store_path, config, self._context, self._given_pipeline_class()) return derived def with_metadata( @@ -7287,15 +7561,17 @@ def with_metadata( metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, ) -> Self: """This binding with `metadata` in place of its own. Nothing is stored; - a legacy `metadata` object or a metadata dict is read as the - constructor reads it, and a legacy object is copied.""" + a legacy `metadata` object or a metadata dict is read as the constructor + reads it, under this binding's context when there is one (so the scope + checks its names), and a legacy object is never shared with the caller.""" source: ArrayMetadataModel | ArrayMetadata if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): - source = _copied_legacy(parse_array_metadata(metadata, str(self._store_path))) + resolved = _legacy_resolved_in_context(metadata, str(self._store_path), self._context) + source = _copied_legacy(resolved) if resolved is metadata else resolved elif isinstance(metadata, _model_classes()): source = _checked_model(metadata) elif isinstance(metadata, Mapping): - source = _bound_metadata_dict(dict(metadata), self._store_path) + source = _bound_metadata_dict(dict(metadata), self._store_path, self._context) else: raise TypeError( "Expected array metadata (a zarr_metadata model, ArrayV2Metadata or " @@ -7318,6 +7594,44 @@ def with_config(self, config: ArrayConfigLike) -> Self: new_config = ArrayConfig(**{**self._config.to_dict(), **config}) # type: ignore[arg-type] return self._derive(self._derived_source(), self._store_path, new_config) + def with_context(self, context: Context | None) -> Self: + """This binding with its metadata resolving in `context`; the data path, + codec pipeline included, derives anew under it. A binding built from a + resolved metadata object is re-resolved there, so a stored name outside + the context fails. `None` rebinds under the default registries and the + globally configured pipeline: each codec or data type whose class is not + the default registry's for its name is read again there from its JSON + (a binding keeps no record of what it held before a context). A name the + defaults do not know fails there, as it would on opening the array: to + change only the pipeline of a binding whose scope holds such a name, + give a context that keeps that scope. A binding over a model (one + `open` made) resolves no names until its data path is built, so a stored + name outside the context fails there instead.""" + # A legacy source is copied once it is resolved, unless resolving made it anew. + source = ( + self._source + if isinstance(self._source, ArrayV2Metadata | ArrayV3Metadata) + else self._derived_source() + ) + old = self._context + if old is not None and isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + # Kinds the old context limited and the new one does not are read + # again in the default registries; kinds the new one limits are + # checked through it when the binding is made. + left_data_types = old.data_types is not None and ( + context is None or context.data_types is None + ) + left_codecs = old.codecs is not None and (context is None or context.codecs is None) + if left_data_types or left_codecs: + source = _reread_in_defaults(source, data_types=left_data_types, codecs=left_codecs) + if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + resolved = _legacy_resolved_in_context(source, str(self._store_path), context) + source = _copied_legacy(resolved) if resolved is self._source else resolved + # Made as the other derivations are, without the subclass's `__init__`. + derived = object.__new__(type(self)) + derived._bind(source, self._store_path, self._config, context, _pipeline_class_for(context)) + return derived + async def store_metadata(self, *, ensure_parents: bool = False) -> None: """Store the metadata documents of this array's Zarr format, unconditionally replacing those stored here. @@ -7459,7 +7773,7 @@ async def _set_selection( if product(indexer.shape) > 0: # Chunks are about to be stored under the repaired metadata, so store it # first: every reader of the store then agrees with them. - await _store_repaired_document(self._store_path, metadata) + await _store_repaired_document(self._store_path, metadata, self._context) return await _set_selection( self._store_path, metadata, @@ -7506,11 +7820,16 @@ def __repr__(self) -> str: def __eq__(self, other: object) -> bool: # Semantic equality over the binding: the document (canonically, via the - # model's own equality), the location and the runtime configuration. + # model's own equality), the location, the runtime configuration, and + # what it resolves with (its context and captured pipeline class, which + # without a context is the one the global configuration named when the + # binding was made). `AsyncArray` keeps the field equality it always had. if not isinstance(other, AsyncArrayIO): return NotImplemented if self._store_path != other._store_path or self._config != other._config: return False + if self._context != other._context or self._pipeline_class is not other._pipeline_class: + return False try: return self.metadata == other.metadata except TypeError: @@ -7537,8 +7856,10 @@ def __init__( metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, 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=context) @classmethod def _wrap(cls, async_io: AsyncArrayIO) -> Self: @@ -7553,10 +7874,13 @@ def build_v2( *, name: str | None = None, config: ArrayConfigLike | None = None, + context: Context | None = None, **kwargs: Unpack[ArrayV2DeclarationKwargs], ) -> Self: """Declare a Zarr format 2 array bound to `store`, storing nothing. See `AsyncArrayIO.build_v2`.""" - return cls._wrap(AsyncArrayIO.build_v2(store, name=name, config=config, **kwargs)) + return cls._wrap( + AsyncArrayIO.build_v2(store, name=name, config=config, context=context, **kwargs) + ) @classmethod def build_v3( @@ -7565,10 +7889,13 @@ def build_v3( *, name: str | None = None, config: ArrayConfigLike | None = None, + context: Context | None = None, **kwargs: Unpack[ArrayV3DeclarationKwargs], ) -> Self: """Declare a Zarr format 3 array bound to `store`, storing nothing. See `AsyncArrayIO.build_v3`.""" - return cls._wrap(AsyncArrayIO.build_v3(store, name=name, config=config, **kwargs)) + return cls._wrap( + AsyncArrayIO.build_v3(store, name=name, config=config, context=context, **kwargs) + ) @classmethod def open( @@ -7578,10 +7905,15 @@ def open( name: str | None = None, zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> Self: """Read the stored 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 @@ -7604,6 +7936,12 @@ def config(self) -> ArrayConfig: """The runtime configuration.""" return self._async_io.config + @property + def context(self) -> Context | None: + """The context this binding resolves extension names and its codec pipeline + under, or `None` for the defaults; `with_context` rebinds it.""" + return self._async_io.context + def to_array(self) -> AnyArray: """An `Array` over a copy of this binding: nothing done through the array (a resize, an attribute update) changes this binding.""" @@ -7624,6 +7962,10 @@ def with_config(self, config: ArrayConfigLike) -> Self: """This binding with `config` merged over its own. See `AsyncArrayIO.with_config`.""" return self._wrap(self._async_io.with_config(config)) + def with_context(self, context: Context | None) -> Self: + """This binding with its metadata resolving in `context`. See `AsyncArrayIO.with_context`.""" + return self._wrap(self._async_io.with_context(context)) + def store_metadata(self, *, ensure_parents: bool = False) -> None: """Store the metadata. See `AsyncArrayIO.store_metadata`.""" sync(self._async_io.store_metadata(ensure_parents=ensure_parents)) diff --git a/src/zarr/core/context.py b/src/zarr/core/context.py index 494945d9f0..75d6bf1447 100644 --- a/src/zarr/core/context.py +++ b/src/zarr/core/context.py @@ -7,33 +7,151 @@ 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 +import copy 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 from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType +_DEFAULT_DATA_TYPES = "zarr.dtype.data_type_registry" + + +def _registry_snapshot(registry: DataTypeRegistry) -> DataTypeRegistry: + """A copy of `registry` as it is now, or the default registry itself.""" + from zarr.core.dtype import data_type_registry + + if registry is data_type_registry: + return registry + # A copy of the registry's own class (a subclass may match names its own way), + # with its own contents. + snapshot = copy.copy(registry) + object.__setattr__(snapshot, "contents", dict(registry.contents)) + object.__setattr__(snapshot, "_lazy_load_list", list(registry._lazy_load_list)) + # Loaded now, so the copy compares the same before and after its first use. + snapshot._lazy_load() + return snapshot + + @dataclass(frozen=True, kw_only=True) 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, and a data type a codec names in its + configuration (`cast_value`'s target), resolve through the default + registries today. + Attributes ---------- - data_types : DataTypeRegistry - The data types a read may resolve. + data_types : DataTypeRegistry | None + The data types a read may resolve, or `None` for the default data type + registry. A registry other than the default is copied when the context + is made, so registering a data type in it later does not change the + context (registering one in the context's own copy does); its pending + entry points load then too, so one that fails to load fails there. A registry + subclass that overrides `match_json` changes how a top-level data type + matches; a data type nested in another (a struct's fields) matches + through the registry's contents. + codecs : Mapping[str, type[Codec]] | None + The codec class each codec name resolves to, or `None` for the + default codec registry. Under a closed codec scope a top-level codec + instance must be of the class the scope maps its name to, so the scope + admits no top-level codec outside it. The name is the one the codec + writes in its JSON: a scope that holds a class only under another name + (such as `BytesCodec` under `endian`; it writes `bytes`) admits none of + its instances, as it could not read back the chain it would store. List + the class under each name it reads and writes. Each class's `from_dict` + must return an instance of that class, which is what the scope binds, and + so must its `evolve_from_array_spec` (returning the instance unchanged once + its JSON is complete): a class that fitting replaces is not held to the + scope. The scope governs Zarr format 3 codecs + only: a Zarr format 2 array's filters and compressor are numcodecs + objects, which no context governs. + codec_pipeline : type[CodecPipeline] | None + The codec pipeline class arrays are read and written with, or `None` + for the pipeline the global configuration selects. A pipeline reads + its own options (such as zarrs-python's `codec_pipeline.strict`) from + the global configuration when it is built; to fix them per binding, + give a subclass that builds under them. + + The data types and codecs in scope decide which names a binding admits and + the classes `metadata_legacy` resolves them to. A pipeline that encodes + from the metadata's JSON with implementations of its own (such as + zarrs-python) runs its own implementation for each admitted name. + + Contexts compare by their fields. A context that sets `codecs` or + `data_types` is not hashable (it holds a mapping or a registry); one that + sets neither is. """ - data_types: DataTypeRegistry + data_types: DataTypeRegistry | None = None + codecs: Mapping[str, type[Codec]] | None = None + codec_pipeline: type[CodecPipeline] | None = None + + def __post_init__(self) -> None: + # Copies, read-only for codecs: the scope must not change under a binding + # that captured it. The default data type registry is kept by reference + # (see `default`). Contexts compare by their fields; one that holds a + # codec mapping or a data type registry is not hashable. + if self.codecs is not None: + object.__setattr__(self, "codecs", MappingProxyType(dict(self.codecs))) + if self.data_types is not None: + object.__setattr__(self, "data_types", _registry_snapshot(self.data_types)) + + def __getstate__(self) -> dict[str, object]: + # A mapping proxy does not pickle; carry the scope as a plain dict. The + # default data type registry is carried by reference, so it unpickles + # as the receiving process's registry rather than a snapshot of ours; + # any other registry unpickles as a copy. + from zarr.core.dtype import data_type_registry + + state = dict(self.__dict__) + if self.codecs is not None: + state["codecs"] = dict(self.codecs) + if self.data_types is data_type_registry: + state["data_types"] = _DEFAULT_DATA_TYPES + 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] + if state.get("data_types") == _DEFAULT_DATA_TYPES: + from zarr.core.dtype import data_type_registry + + object.__setattr__(self, "data_types", data_type_registry) + + @classmethod + def _over(cls, data_types: DataTypeRegistry) -> Context: + """A context over `data_types` itself, not a copy: for a resolver that lives + only as long as one call on that registry.""" + context = object.__new__(cls) + object.__setattr__(context, "data_types", data_types) + object.__setattr__(context, "codecs", None) + object.__setattr__(context, "codec_pipeline", None) + return context @classmethod def default(cls) -> Context: @@ -41,12 +159,49 @@ def default(cls) -> Context: The context of the default registries, such as `zarr.dtype.data_type_registry`. It refers to the registries themselves, so it includes extensions registered after it is - created. + created; it is `Context()`, which leaves every kind unset. """ - # avoid circular import - from zarr.core.dtype import data_type_registry + return cls() + + @property + def limits_names(self) -> bool: + """Whether this context limits which names resolve: it sets `data_types` + or `codecs`. A context that only selects the codec pipeline resolves + exactly as no context does.""" + return self.data_types is not None or self.codecs is not None + + def data_type_registry(self) -> DataTypeRegistry: + """The data types this context resolves: its own, or the default registry.""" + if self.data_types is None: + from zarr.core.dtype import data_type_registry + + return data_type_registry + return self.data_types + + 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 cls(data_types=data_type_registry) + return get_pipeline_class() + return self.codec_pipeline type JSONLocation = tuple[str | int, ...] @@ -100,4 +255,4 @@ def resolve_data_type(self, data: DTypeJSON) -> ZDType[TBaseDType, TBaseScalar]: ValueError If no data type in the context matches `data`. """ - return self.context.data_types._match_json(data, resolver=self) + return self.context.data_type_registry()._match_json(data, resolver=self) diff --git a/src/zarr/core/dtype/__init__.py b/src/zarr/core/dtype/__init__.py index d1dbd6e2c8..4ef0bba6e0 100644 --- a/src/zarr/core/dtype/__init__.py +++ b/src/zarr/core/dtype/__init__.py @@ -1,8 +1,10 @@ from __future__ import annotations from collections.abc import Sequence -from typing import TYPE_CHECKING, Final +from dataclasses import replace +from typing import TYPE_CHECKING, Any, Final, cast +from zarr.core.dtype.common import _UNSTABLE_WARNING_SILENCED, HasEndianness from zarr.core.dtype.npy.bool import Bool from zarr.core.dtype.npy.bytes import ( NullTerminatedBytes, @@ -35,6 +37,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 +164,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 +182,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_type_registry() + 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_type_registry() + return registry.match_json(dtype_spec, zarr_format=zarr_format) def parse_data_type( @@ -234,6 +247,7 @@ def parse_dtype( dtype_spec: ZDTypeLike, *, zarr_format: ZarrFormat, + context: Context | None = None, ) -> ZDType[TBaseDType, TBaseScalar]: """ Convert the input as a ZDType. @@ -268,21 +282,44 @@ def parse_dtype( ``` """ if isinstance(dtype_spec, ZDType): + if context is None or context.data_types is None: + return dtype_spec + # A closed world checks instances too: the scope must read the instance's + # JSON as the instance's own class, as a closed codec scope requires. + matched = _matched_in_scope(dtype_spec, context, zarr_format) + if not _same_classes(matched, dtype_spec): + raise ValueError( + f"Data type {dtype_spec!r} is not of the classes this context's scope " + f"reads its JSON as, {matched!r}. Give an instance of those classes, or " + "the data type's name." + ) return dtype_spec # 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 and context.data_types 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 _matched_in_scope( + cast("ZDType[TBaseDType, TBaseScalar]", utf8), context, 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] + zdtype = get_data_type_from_native_dtype(dtype_spec, context=context) # type: ignore[arg-type] + if context is not None and context.data_types is not None: + # A native data type can hold others (a structured dtype's fields), which + # resolve from the defaults: check the whole through the closed world. + return _matched_in_scope(zdtype, context, zarr_format) + return zdtype def __getattr__(name: str) -> object: @@ -299,3 +336,66 @@ def __getattr__(name: str) -> object: ) return DataTypeValidationError raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def _json_to_reread(zdtype: ZDType[TBaseDType, TBaseScalar], zarr_format: ZarrFormat) -> DTypeJSON: + """`zdtype`'s JSON, for reading it again: without the warning a data type with + no stable specification gives when it is written (nothing is written). Only the + warning given through `v3_unstable_dtype_warning` is silenced; a data type that + calls `warnings.warn` itself warns here as well.""" + token = _UNSTABLE_WARNING_SILENCED.set(True) + try: + return zdtype.to_json(zarr_format=zarr_format) + finally: + _UNSTABLE_WARNING_SILENCED.reset(token) + + +def _matched_in_scope( + zdtype: ZDType[TBaseDType, TBaseScalar], context: Context, zarr_format: ZarrFormat +) -> ZDType[TBaseDType, TBaseScalar]: + """The data type a closed data type scope matches `zdtype`'s JSON to, in the byte + order of `zdtype` (see `_with_byte_order_of`): how a resolved data type is read + again under a scope, which rebinds the scope's class for its name.""" + matched = context.data_type_registry().match_json( + _json_to_reread(zdtype, zarr_format), zarr_format=zarr_format + ) + return _with_byte_order_of(matched, zdtype) + + +def _same_classes( + zdtype: ZDType[TBaseDType, TBaseScalar], other: ZDType[TBaseDType, TBaseScalar] +) -> bool: + """Whether `zdtype` and `other` are of the same class, and so are their fields, for + a structured data type.""" + if type(zdtype) is not type(other): + return False + if isinstance(zdtype, Structured) and isinstance(other, Structured): + return len(zdtype.fields) == len(other.fields) and all( + _same_classes(field, other_field) + for (_, field), (_, other_field) in zip(zdtype.fields, other.fields, strict=True) + ) + return True + + +def _with_byte_order_of( + zdtype: ZDType[TBaseDType, TBaseScalar], source: ZDType[TBaseDType, TBaseScalar] +) -> ZDType[TBaseDType, TBaseScalar]: + """`zdtype`, read again from `source`'s JSON, in `source`'s byte order (its fields', + for a structured data type): Zarr format 3 data type JSON carries none, so reading + it alone gives the host's byte order.""" + if ( + isinstance(source, Structured) + and isinstance(zdtype, Structured) + and len(source.fields) == len(zdtype.fields) + ): + fields = tuple( + (name, _with_byte_order_of(field, source_field)) + for (name, field), (_, source_field) in zip(zdtype.fields, source.fields, strict=True) + ) + return cast("ZDType[TBaseDType, TBaseScalar]", replace(cast("Any", zdtype), fields=fields)) + if isinstance(source, HasEndianness) and isinstance(zdtype, HasEndianness): + return cast( + "ZDType[TBaseDType, TBaseScalar]", + replace(cast("Any", zdtype), endianness=source.endianness), + ) + return zdtype diff --git a/src/zarr/core/dtype/common.py b/src/zarr/core/dtype/common.py index 342c0d09d5..42c5d1f681 100644 --- a/src/zarr/core/dtype/common.py +++ b/src/zarr/core/dtype/common.py @@ -1,5 +1,6 @@ from __future__ import annotations +import contextvars import sys import warnings from collections.abc import Mapping, Sequence @@ -246,10 +247,20 @@ class HasObjectCodec: object_codec_id: ClassVar[str] +_UNSTABLE_WARNING_SILENCED: contextvars.ContextVar[bool] = contextvars.ContextVar( + "_UNSTABLE_WARNING_SILENCED", default=False +) +"""Set while zarr reads a data type's JSON back without writing it, so the warning, +which is about what is written, is not given. Per thread and task, unlike the global +warning filters.""" + + def v3_unstable_dtype_warning(dtype: object) -> None: """ Emit this warning when a data type does not have a stable zarr v3 spec """ + if _UNSTABLE_WARNING_SILENCED.get(): + return msg = ( f"The data type ({dtype}) does not have a Zarr V3 specification. " "That means that the representation of arrays saved with this data type may change without " diff --git a/src/zarr/core/dtype/registry.py b/src/zarr/core/dtype/registry.py index 1fdd97e7af..9eac6d55e6 100644 --- a/src/zarr/core/dtype/registry.py +++ b/src/zarr/core/dtype/registry.py @@ -270,7 +270,7 @@ def match_json( If no matching Zarr data type is found for the given JSON data. """ - resolver = Resolver(context=Context(data_types=self), zarr_format=zarr_format) + resolver = Resolver(context=Context._over(self), zarr_format=zarr_format) return self._match_json(data, resolver=resolver) def _match_json( diff --git a/src/zarr/core/metadata/io.py b/src/zarr/core/metadata/io.py index 024c9a57bb..ab0b71dfe7 100644 --- a/src/zarr/core/metadata/io.py +++ b/src/zarr/core/metadata/io.py @@ -20,6 +20,7 @@ from zarr.core.buffer import Buffer from zarr.core.common import JSON, ZarrFormat + from zarr.core.context import Context from zarr.core.group import GroupMetadata from zarr.core.metadata import ArrayMetadata from zarr.core.metadata.repair import ArrayDocument @@ -149,12 +150,15 @@ def reads_as(document: ArrayDocument, metadata: ArrayMetadata) -> bool: return read_repaired(document, metadata.zarr_format).to_dict() == metadata.to_dict() -def read_repaired(document: ArrayDocument, zarr_format: ZarrFormat) -> ArrayMetadata: - """The metadata the stored array document `document` reads as, without warnings.""" +def read_repaired( + document: ArrayDocument, zarr_format: ZarrFormat, context: Context | None = None +) -> ArrayMetadata: + """The metadata the stored array document `document` reads as, without warnings, + resolving its names in `context`.""" from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata reader = ArrayV2Metadata if zarr_format == 2 else ArrayV3Metadata - metadata, _ = reader._read_stored(dict(document)) + metadata, _ = reader._read_stored(dict(document), context=context) return metadata @@ -218,7 +222,11 @@ async def upsert_metadata( def parse_stored_array( - documents: Mapping[str, Buffer], zarr_format: ZarrFormat, path: str | None = None + documents: Mapping[str, Buffer], + zarr_format: ZarrFormat, + path: str | None = None, + *, + context: Context | None = None, ) -> ArrayMetadata: """The metadata of an array from its documents (by store key, see `ARRAY_DOCUMENTS`), read with the repairs but without their warnings (whoever asks has warned, or reads @@ -241,7 +249,7 @@ def parse_stored_array( else: raise ArrayNotFoundError(f"No Zarr format {zarr_format} array metadata document.") repaired, _ = repair_array_document(stored, zarr_format) - return parse_array_metadata(dict(repaired), path) + return parse_array_metadata(dict(repaired), path, context=context) def _build_parents(store_path: StorePath, zarr_format: ZarrFormat) -> dict[str, GroupMetadata]: diff --git a/src/zarr/core/metadata/v2.py b/src/zarr/core/metadata/v2.py index 6b32b41676..1de2839c96 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, @@ -149,16 +150,22 @@ 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`.""" - metadata, readings = cls._read_stored(data) + 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.""" + metadata, readings = cls._read_stored(data, context=context) warn_repairs(readings, path) return metadata @classmethod - def _read_stored(cls, data: dict[str, Any]) -> tuple[ArrayV2Metadata, list[Reading]]: + def _read_stored( + cls, data: dict[str, Any], *, context: Context | None = None + ) -> tuple[ArrayV2Metadata, list[Reading]]: """`from_dict` without its warnings: the metadata, and how the repairs read it.""" repaired, readings = repair_array_document(data, 2) # a new dict, because we are modifying it @@ -182,7 +189,7 @@ def _read_stored(cls, data: dict[str, Any]) -> tuple[ArrayV2Metadata, list[Readi "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 186e68f1c3..d7127e130f 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -44,7 +44,7 @@ repair_array_document, warn_repairs, ) -from zarr.errors import MetadataValidationError, NodeTypeValidationError +from zarr.errors import MetadataValidationError, NodeTypeValidationError, UnknownCodecError from zarr.registry import get_codec_class if TYPE_CHECKING: @@ -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 Reading @@ -70,7 +71,13 @@ def parse_node_type_array(data: object) -> Literal["array"]: ) -def parse_codecs(data: object) -> tuple[Codec, ...]: +def _written_name(codec: Codec) -> str: + """The name `codec` writes in its JSON, which a codec scope looks it up by.""" + name, _ = parse_named_configuration(codec.to_dict(), require_configuration=False) + return name + + +def parse_codecs(data: object, *, context: Context | None = None) -> tuple[Codec, ...]: out: tuple[Codec, ...] = () if not isinstance(data, Iterable): @@ -80,11 +87,30 @@ def parse_codecs(data: object) -> tuple[Codec, ...]: if isinstance( c, ArrayArrayCodec | ArrayBytesCodec | BytesBytesCodec ): # Can't use Codec here because of mypy limitation - out += (c,) + if context is None or context.codecs is None: + out += (c,) + else: + # A closed world checks instances too: the instance's name must + # be in it, and the instance must be of the class it maps the name + # to (rebuilding one would drop state its JSON does not hold, such + # as a codec's attributes still to be fitted to the data type). + name = _written_name(c) + scoped_class = context.codec_class(name) + if type(c) is not scoped_class: + raise UnknownCodecError( + f"Codec {name!r} is a {type(c).__name__} instance; this context's " + f"scope maps {name!r} to {scoped_class.__name__}. Give a " + f"{scoped_class.__name__} instance, or the codec's JSON." + ) + out += (c,) 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: @@ -99,6 +125,13 @@ def parse_codecs(data: object) -> tuple[Codec, ...]: raise MetadataValidationError( f"KeyError{key_text} while parsing the configuration for codec {name_parsed!r}." ) from e + except TypeError as e: + # A closed scope's class refusing the JSON reports as it does when + # already resolved metadata enters the scope (`_codec_in_scope`). + # Caught after `KeyError`, so an error that is both converts as before. + if context is None or context.codecs is None: + raise + raise _refused_by_scope(name_parsed, codec_cls, c, e) from e return out @@ -662,11 +695,16 @@ 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; under a + context that limits codecs they resolve before the other fields are + checked, so a document with several invalid fields reports a codec first.""" # 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": @@ -676,12 +714,14 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self: if path is not None: e.add_note(f"Array {path!r}: nothing was read.") raise - metadata, readings = cls._read_stored(data) + metadata, readings = cls._read_stored(data, context=context) warn_repairs(readings, path) return metadata @classmethod - def _read_stored(cls, data: dict[str, JSON]) -> tuple[Self, list[Reading]]: + def _read_stored( + cls, data: dict[str, JSON], *, context: Context | None = None + ) -> tuple[Self, list[Reading]]: """`from_dict` without its warnings or the rectilinear chunks flag: the metadata, and how the repairs read it.""" repaired, readings = repair_array_document(data, 3) @@ -696,7 +736,7 @@ def _read_stored(cls, data: dict[str, JSON]) -> tuple[Self, list[Reading]]: 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: @@ -718,7 +758,15 @@ def _read_stored(cls, data: dict[str, JSON]) -> tuple[Self, list[Reading]]: 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"], + # Under a context, codec names resolve here, through it; `__init__`'s + # own `parse_codecs` passes the resolved instances through untouched. + # Without one, `__init__` parses them, after the other fields, as it + # always has. + codecs=( + _data_typed["codecs"] + if context is None or context.codecs is None + else 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] @@ -728,7 +776,7 @@ def _read_stored(cls, data: dict[str, JSON]) -> tuple[Self, list[Reading]]: extra_fields=allowed_extra_fields, storage_transformers=_data_typed.get("storage_transformers", ()), # type: ignore[arg-type] ) - return metadata, readings + return cast("Self", _rescoped(metadata, context)), readings def to_dict(self) -> dict[str, JSON]: out_dict = super().to_dict() @@ -778,3 +826,63 @@ def update_shape(self, shape: tuple[int, ...]) -> Self: def update_attributes(self, attributes: dict[str, JSON]) -> Self: return replace(self, attributes=attributes) + + +def _codec_in_scope(codec: Codec, context: Context) -> Codec: + """`codec`, of metadata already resolved, read again by name under `context`: the + instance itself if it is of the class `context` maps its name to (with no codec + scope, the default registry's), else that class read from its JSON, which + carries only the state the JSON holds. A builder refuses such an instance + instead, as a codec it is given may not be fitted to the array yet.""" + name = _written_name(codec) + try: + scoped_class = context.codec_class(name) + except UnknownCodecError as e: + e.add_note(f"{name!r} is the name a {type(codec).__name__} writes.") + raise + if type(codec) is scoped_class: + return codec + spec = codec.to_dict() + try: + rebuilt = scoped_class.from_dict(spec) + except (TypeError, KeyError) as e: + raise _refused_by_scope(name, scoped_class, spec, e) from e + # What it built must read back here: of the class its own written name maps to. + written = _written_name(rebuilt) + try: + written_class = context.codec_class(written) + except UnknownCodecError as e: + e.add_note( + f"{written!r} is the name a {type(rebuilt).__name__} writes, read from {name!r}." + ) + raise + if written_class is not type(rebuilt): + raise UnknownCodecError( + f"Codec {name!r} resolves to {scoped_class.__name__}, whose from_dict built " + f"a {type(rebuilt).__name__} writing {written!r}, which resolves to " + f"{written_class.__name__} here." + ) + return rebuilt + + +def _refused_by_scope( + name: str, scoped_class: type[Codec], spec: object, error: Exception +) -> MetadataValidationError: + """The error for a codec whose JSON the class its name resolves to refuses: the + same whichever route read the codec.""" + return MetadataValidationError( + f"Codec {name!r} resolves to {scoped_class.__name__}, which refuses {spec!r} ({error!r})." + ) + + +def _rescoped(metadata: ArrayV3Metadata, context: Context | None) -> ArrayV3Metadata: + """`metadata` with the codecs that fitting to the array rebuilt as another class + than a closed codec scope's (the numcodecs wrappers rebuild as their own base + class) read again as the scope's class, from their JSON, which fitting has + completed; itself when there are none.""" + if context is None or context.codecs is None: + return metadata + codecs = tuple(_codec_in_scope(codec, context) for codec in metadata.codecs) + if all(new is old for new, old in zip(codecs, metadata.codecs, strict=True)): + return metadata + return replace(metadata, codecs=codecs) 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/tests/test_array.py b/tests/test_array.py index 40bde50783..be2610edf5 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -56,7 +56,7 @@ resolve_outer_and_inner_chunks, ) from zarr.core.chunk_key_encodings import ChunkKeyEncodingParams -from zarr.core.codec_pipeline import BatchedCodecPipeline +from zarr.core.codec_pipeline import BatchedCodecPipeline, FusedCodecPipeline from zarr.core.common import JSON, ChunksLike, ZarrFormat, ceildiv from zarr.core.dtype import ( DateTime64, @@ -82,6 +82,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 @@ -2172,20 +2173,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(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: diff --git a/tests/test_array_io.py b/tests/test_array_io.py index c7f357fc68..26281b69f3 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -8,7 +8,9 @@ import math import pickle import threading +import warnings from dataclasses import MISSING, dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Any, SupportsIndex, cast from unittest import mock @@ -20,20 +22,43 @@ import zarr from zarr import ArrayIO, AsyncArray, AsyncArrayIO from zarr.api.asynchronous import open_array -from zarr.core.array import create_codec_pipeline, parse_array_metadata +from zarr.codecs import BloscCodec, BytesCodec, GzipCodec, ShardingCodec, ZstdCodec +from zarr.codecs.numcodecs import Shuffle +from zarr.core.array import build_array_metadata_v3, create_codec_pipeline, parse_array_metadata +from zarr.core.array_spec import parse_array_config from zarr.core.buffer import default_buffer_prototype from zarr.core.codec_pipeline import BatchedCodecPipeline, FusedCodecPipeline +from zarr.core.dtype import ( + Float32, + Int8, + NullTerminatedBytes, + Struct, + UInt16, + VariableLengthUTF8, + data_type_registry, + parse_dtype, +) +from zarr.core.dtype.registry import DataTypeRegistry from zarr.core.metadata import ArrayMetadata, ArrayV2Metadata, ArrayV3Metadata from zarr.core.metadata.stored import stored_documents from zarr.core.sync import sync from zarr.dtype import Int16 -from zarr.errors import ArrayNotFoundError, UnknownCodecError, ZarrUserWarning +from zarr.errors import ( + ArrayNotFoundError, + UnknownCodecError, + UnstableSpecificationWarning, + ZarrUserWarning, +) +from zarr.errors import MetadataValidationError as ZarrMetadataValidationError +from zarr.experimental import Context from zarr.storage import LocalStore, MemoryStore, StorePath if TYPE_CHECKING: from collections.abc import Callable, Mapping from pathlib import Path + from typing import Self + from zarr.core.array_spec import ArrayConfigLike from zarr.core.common import JSON, ZarrFormat METADATA_DOCUMENTS = ("zarr.json", ".zarray", ".zattrs", ".zgroup") @@ -749,8 +774,6 @@ async def test_erase_metadata_removes_array_documents_however_written( async def test_with_metadata_takes_any_mapping() -> None: - from types import MappingProxyType - aio = AsyncArrayIO.build_v3( MemoryStore(), name="foo", shape=(4,), data_type="uint8", chunk_grid=(4,) ) @@ -759,13 +782,659 @@ async def test_with_metadata_takes_any_mapping() -> None: await rebound.store_metadata() +BYTES_JSON: dict[str, JSON] = {"name": "bytes", "configuration": {"endian": "little"}} + + +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.""" + + +async def test_context_scopes_resolution() -> None: + """Names resolve through the context when the data path is built, the + context's pipeline wins over the global choice, and with_context rebinds.""" + 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 = AsyncArrayIO.build_v3( + store, + name="foo", + shape=data.shape, + data_type="uint16", + chunk_grid=(2, 2), + codecs=[BYTES_JSON, 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) + + # the context survives derivation on both layers + assert isinstance(opened.with_config({"order": "F"}).to_array().codec_pipeline, _MarkedPipeline) + assert isinstance(opened.to_array().with_config({"order": "F"}).codec_pipeline, _MarkedPipeline) + + rebound = opened.with_context(None) + assert not isinstance(rebound.to_array().codec_pipeline, _MarkedPipeline) + np.testing.assert_array_equal(await rebound.getitem(Ellipsis), data) + + +async def test_context_unknown_codec_resolution_raises() -> None: + """A stored codec name outside the context opens for metadata work and + fails to resolve when the data path is built, naming the scope.""" + store = MemoryStore() + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(4,)) + 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() # metadata work needs no resolution + with pytest.raises(UnknownCodecError, match="Codecs in scope: \\['bytes'\\]"): + await opened.getitem(Ellipsis) + + +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"): + AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[BYTES_JSON, ZSTD_JSON], + context=context, + ) + + +def test_context_unknown_dtype_raises() -> None: + context = Context(data_types=DataTypeRegistry()) + with pytest.raises(ValueError, match="No Zarr data type found"): + AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint16", chunk_grid=(4,), context=context + ) + # the str shortcut goes through the registry too + with pytest.raises(ValueError, match="No Zarr data type found"): + AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type=str, chunk_grid=(4,), context=context + ) + + +def test_context_pickles_with_codecs() -> None: + """An array bound to a context with a codec scope pickles and unpickles + with the scope intact.""" + context = Context( + data_types=data_type_registry, + codecs={"bytes": BytesCodec, "zstd": ZstdCodec}, + codec_pipeline=_MarkedPipeline, + ) + arr = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(4,)) + rebound = arr.io.with_context(context).to_array() + revived = pickle.loads(pickle.dumps(rebound)) + revived_context = revived._async_array._io._context + assert revived_context is not None + assert revived_context.codecs is not None + assert set(revived_context.codecs) == {"bytes", "zstd"} + assert revived_context.codec_pipeline is _MarkedPipeline + np.testing.assert_array_equal(revived[:], np.zeros(4, dtype="uint8")) + + +async def test_with_context_reresolves_a_legacy_source() -> None: + """A binding built from resolved metadata is re-resolved by with_context, + so its stored names are checked against the scope.""" + arr = await zarr.api.asynchronous.create_array( + MemoryStore(), shape=(4,), dtype="uint16", chunks=(4,) + ) + full = Context( + data_types=data_type_registry, + codecs={"bytes": BytesCodec, "zstd": ZstdCodec}, + codec_pipeline=_MarkedPipeline, + ) + rebound = arr.io.with_context(full) + assert isinstance(rebound.to_array().codec_pipeline, _MarkedPipeline) + data = np.arange(4, dtype="uint16") + await rebound.setitem(Ellipsis, data) + np.testing.assert_array_equal(await rebound.getitem(Ellipsis), data) + + +async def test_with_context_rejects_a_legacy_codec_outside_the_scope() -> None: + arr = await zarr.api.asynchronous.create_array( + MemoryStore(), shape=(4,), dtype="uint16", chunks=(4,) + ) + narrow = Context(data_types=data_type_registry, codecs={"bytes": BytesCodec}) + with pytest.raises(UnknownCodecError, match="is not in this context"): + arr.io.with_context(narrow) + + +async def test_with_context_rejects_a_legacy_dtype_outside_the_scope() -> None: + arr = await zarr.api.asynchronous.create_array( + MemoryStore(), shape=(4,), dtype="uint16", chunks=(4,) + ) + with pytest.raises(ValueError, match="No Zarr data type found"): + arr.io.with_context(Context(data_types=DataTypeRegistry())) + + +async def test_context_v2_scopes_dtype_resolution() -> None: + """A v2 document resolves its dtype through the context on build and open.""" + store = MemoryStore() + registry = DataTypeRegistry() + registry.register(UInt16._zarr_v3_name, UInt16) + context = Context(data_types=registry) + data = np.arange(4, dtype="uint16") + aio = AsyncArrayIO.build_v2( + store, name="foo", shape=(4,), dtype="uint16", chunks=(4,), context=context + ) + await aio.store_metadata() + await aio.setitem(Ellipsis, data) + opened = await AsyncArrayIO.open(store, name="foo", context=context) + np.testing.assert_array_equal(await opened.getitem(Ellipsis), data) + + +def test_context_v2_unknown_dtype_raises() -> None: + context = Context(data_types=DataTypeRegistry()) + with pytest.raises(ValueError, match="No Zarr data type found"): + AsyncArrayIO.build_v2( + MemoryStore(), name="foo", shape=(4,), dtype="uint16", chunks=(4,), context=context + ) + + +def test_context_default_codecs_outside_the_scope_raise() -> None: + """The default codec chain resolves under the context too, so a closed + scope never stores a chain it could not read back.""" + narrow = Context(data_types=data_type_registry, codecs={"bytes": BytesCodec}) + with pytest.raises(UnknownCodecError, match="is not in this context"): + AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint16", chunk_grid=(4,), context=narrow + ) + + +async def test_with_context_rebinding_stores_the_repair() -> None: + """Rebinding a repaired legacy array carries nothing about the store: its + store's cache still records the broken document, so the first write + through the rebound binding stores the repair.""" + store = MemoryStore() + doc: dict[str, Any] = { + "zarr_format": 2, + "shape": [3], + "chunks": [0], + "dtype": " None: + """The variable-length-UTF8 shortcut returns what the context's registry + matches, not the default class.""" + registry = DataTypeRegistry() + registry.register(VariableLengthUTF8._zarr_v3_name, _ScopedUTF8) # type: ignore[arg-type] + context = Context(data_types=registry) + assert isinstance(parse_dtype(str, zarr_format=2, context=context), _ScopedUTF8) + assert isinstance(parse_dtype(str, zarr_format=3, context=context), _ScopedUTF8) + + +def test_with_metadata_under_a_context_rejects_a_codec_outside_it() -> None: + narrow = Context(data_types=data_type_registry, codecs={"bytes": BytesCodec}) + aio = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[BYTES_JSON], + context=narrow, + ) + outside = zarr.create_array( + MemoryStore(), shape=(4,), dtype="uint8", chunks=(4,) + ).io.metadata_legacy + with pytest.raises(UnknownCodecError, match="is not in this context"): + aio.with_metadata(outside) + + +def test_builders_under_a_context_reject_instances_outside_it() -> None: + """Data type and codec instances are checked against a closed context + too, so a builder never binds what its own context cannot read back.""" + registry = DataTypeRegistry() + registry.register(UInt16._zarr_v3_name, UInt16) + context = Context(data_types=registry, codecs={"bytes": BytesCodec, "zstd": ZstdCodec}) + with pytest.raises(ValueError, match="No Zarr data type found"): + AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type=Int16(), chunk_grid=(4,), context=context + ) + with pytest.raises(UnknownCodecError, match="'gzip' is not in this context"): + AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint16", + chunk_grid=(4,), + codecs=[BytesCodec(), GzipCodec()], + context=context, + ) + + +class _ScopedUInt16(UInt16): + """A uint16 data type distinguishable from the default.""" + + +def _scoped_uint16_registry() -> DataTypeRegistry: + registry = DataTypeRegistry() + registry.register(UInt16._zarr_v3_name, _ScopedUInt16) + return registry + + +def test_v2_bindings_resolve_through_a_data_type_scope() -> None: + """A Zarr format 2 legacy object entering a data type scope (constructor, + with_metadata) binds the scope's class, and leaving it rebinds the default; + a codec scope does not govern a Zarr format 2 array's numcodecs.""" + context = Context(data_types=_scoped_uint16_registry()) + legacy = zarr.create_array( + MemoryStore(), shape=(4,), dtype="uint16", chunks=(4,), zarr_format=2 + ).io.metadata_legacy + declared = AsyncArrayIO.build_v2( + MemoryStore(), shape=(4,), dtype="uint16", chunks=(4,), context=context + ) + for binding in ( + AsyncArrayIO(legacy, StorePath(MemoryStore()), context=context), + declared.with_metadata(legacy), + ): + assert type(binding.metadata_legacy.dtype) is _ScopedUInt16 + assert type(binding.with_context(None).metadata_legacy.dtype) is UInt16 + + codecs_only = Context(codecs={"bytes": BytesCodec}) + unscoped = AsyncArrayIO(legacy, StorePath(MemoryStore()), context=codecs_only) + assert unscoped.metadata_legacy.to_dict() == legacy.to_dict() + + +def test_switching_between_scopes_of_different_kinds() -> None: + """Moving from a codec scope to a data type scope reads the codecs again + under the defaults and resolves the data type through the new scope.""" + codec_scoped = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint16", + chunk_grid=(4,), + codecs=[BYTES_JSON], + context=Context(codecs={"bytes": _UnregisteredBytes}), + ) + assert type(codec_scoped.metadata_legacy.codecs[0]) is _UnregisteredBytes # type: ignore[union-attr] + switched = codec_scoped.with_context(Context(data_types=_scoped_uint16_registry())) + assert type(switched.metadata_legacy.codecs[0]) is BytesCodec # type: ignore[union-attr] + assert type(switched.metadata_legacy.dtype) is _ScopedUInt16 + + +def test_bindings_resolve_through_the_context_they_are_under() -> None: + """Whatever reaches a binding under a context, its data type is the + context's class: a builder's declaration, a chunk-size-0 v2 declaration + (read as repaired), and a legacy object given to with_metadata; leaving + the context rebinds the default class.""" + context = Context(data_types=_scoped_uint16_registry()) + + built = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint16", chunk_grid=(4,), context=context + ) + zero_chunk = AsyncArrayIO.build_v2( + MemoryStore(), shape=(3,), dtype="uint16", chunks=(0,), context=context + ) + outside = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint16", chunks=(4,)) + rebound = built.with_metadata(outside.io.metadata_legacy) + + for binding in (built, zero_chunk, rebound): + assert isinstance(binding.metadata_legacy.dtype, _ScopedUInt16) + assert zero_chunk.metadata_legacy.chunks == (1,) + assert type(built.with_context(None).metadata_legacy.dtype) is UInt16 + + +def test_context_selects_a_pipeline_alone() -> None: + """A context that only selects the codec pipeline resolves names from the + default registries.""" + context = Context(codec_pipeline=_MarkedPipeline) + aio = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint16", chunk_grid=(4,), context=context + ) + assert isinstance(aio.to_array().codec_pipeline, _MarkedPipeline) + assert type(aio.metadata_legacy.dtype) is UInt16 + + +def test_bindings_under_different_contexts_differ() -> None: + store = MemoryStore() + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(4,)) + scoped = aio.with_context(Context(codec_pipeline=_MarkedPipeline)) + assert scoped != aio + assert scoped.with_context(None) == aio + + +class _UnregisteredUInt16(UInt16): + """A data type no registry holds.""" + + +class _UnregisteredBytes(BytesCodec): + """A codec no registry holds by its own name; it encodes as `bytes`.""" + + +def test_pipeline_only_context_resolves_as_no_context() -> None: + """A context that only selects the codec pipeline accepts exactly what no + context accepts: unregistered data type and codec instances bind as given, + for both formats, and survive with_config and with_context.""" + context = Context(codec_pipeline=_MarkedPipeline) + v3 = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type=_UnregisteredUInt16(), + chunk_grid=(4,), + codecs=[_UnregisteredBytes()], + context=context, + ) + v2 = AsyncArrayIO.build_v2( + MemoryStore(), shape=(4,), dtype=_UnregisteredUInt16(), chunks=(4,), context=context + ) + v3_bindings = [ + v3, + v3.with_config({"order": "F"}), + v3.with_context(None), + AsyncArrayIO(v3.metadata_legacy, StorePath(MemoryStore()), context=context), + v3.with_metadata(v3.metadata_legacy), + ] + for binding in [*v3_bindings, v2]: + assert type(binding.metadata_legacy.dtype) is _UnregisteredUInt16 + for binding in v3_bindings: + assert type(binding.metadata_legacy.codecs[0]) is _UnregisteredBytes # type: ignore[union-attr] + assert isinstance(v3.to_array().codec_pipeline, _MarkedPipeline) + + +async def test_scoped_array_keeps_its_context() -> None: + """An array under a closed context keeps it across with_config and + pickling, reads chunk-size-0 v2 metadata as repaired without a warning, + and binds the context's classes.""" + registry = DataTypeRegistry() + registry.register(UInt16._zarr_v3_name, _ScopedUInt16) + context = Context(data_types=registry, codec_pipeline=_MarkedPipeline) + zero_chunk = ArrayV2Metadata(shape=(3,), chunks=(0,), dtype=UInt16(), fill_value=0, order="C") + store = MemoryStore() + with warnings.catch_warnings(): + warnings.simplefilter("error") + arr = AsyncArray(zero_chunk, StorePath(store), context=context) + declared = AsyncArrayIO.build_v2(store, shape=(3,), dtype="uint16", chunks=(3,)) + rebound = declared.with_context(context).with_metadata(zero_chunk) + configured = arr.with_config({"order": "F"}) + revived = pickle.loads(pickle.dumps(arr)) + for array in (arr, configured, revived, rebound.to_array()): + assert array.metadata.chunks == (1,) + assert isinstance(array.metadata.dtype, _ScopedUInt16) + assert array.io._context == context + assert isinstance(array.codec_pipeline, _MarkedPipeline) + + +def test_default_context_is_the_empty_context() -> None: + assert Context.default() == Context() + assert not Context().limits_names + assert Context(codecs={"bytes": BytesCodec}).limits_names + + +def test_closed_data_type_scope_checks_structured_fields() -> None: + """A NumPy structured dtype is checked whole: a field type outside a closed + data type scope is rejected at build, as a Struct instance is.""" + registry = DataTypeRegistry() + registry.register(Int8._zarr_v3_name, Int8) + registry.register(Struct._zarr_v3_name, Struct) + context = Context(data_types=registry) + with pytest.raises(ValueError, match="No Zarr data type found"): + AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type=np.dtype([("a", "i1"), ("b", "f8")]), + chunk_grid=(4,), + context=context, + ) + + +def test_closed_codec_scope_keeps_admitted_instances() -> None: + """An admitted codec instance binds as given, so it evolves against the data + type as it does with no context.""" + closed = Context(codecs={"bytes": BytesCodec, "blosc": BloscCodec}) + stored = [ + AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="int32", + chunk_grid=(4,), + codecs=[BytesCodec(), BloscCodec()], + context=context, + ).metadata_legacy.to_dict()["codecs"] + for context in (closed, None) + ] + assert stored[0] == stored[1] + + +def test_leaving_a_closed_scope_rebinds_the_defaults() -> None: + """Rebinding under any context that limits no names (or none) reads the + names a closed scope resolved again under the defaults, and a context that + sets nothing is no context.""" + scoped = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[BYTES_JSON], + context=Context(codecs={"bytes": _UnregisteredBytes}), + ) + for context in (None, Context(), Context(codec_pipeline=_MarkedPipeline)): + rebound = scoped.with_context(context) + assert type(rebound.metadata_legacy.codecs[0]) is BytesCodec # type: ignore[union-attr] + plain = scoped.with_context(None) + assert plain.with_context(Context.default()) == plain + + +def test_with_metadata_under_a_pipeline_only_context_copies() -> None: + declared = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + context=Context(codec_pipeline=_MarkedPipeline), + ) + mine = zarr.create_array( + MemoryStore(), shape=(4,), dtype="uint8", chunks=(4,) + ).io.metadata_legacy + rebound = declared.with_metadata(mine) + assert rebound.metadata_legacy is not mine + assert rebound.metadata_legacy.attributes is not mine.attributes + + +def test_the_context_is_not_a_dataclass_field() -> None: + """The dataclass fields are as they always were; `dataclasses.replace` builds + from them, so its result has no context, and `with_config` keeps it.""" + context = Context(codec_pipeline=_MarkedPipeline) + arr = AsyncArray( + zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(4,)).metadata, + StorePath(MemoryStore()), + context=context, + ) + order_f = parse_array_config({"order": "F"}) + assert dataclasses.replace(arr, config=order_f).io.context is None + configured = arr.with_config(order_f) + assert configured.io.context == context + assert isinstance(configured.codec_pipeline, _MarkedPipeline) + assert [f.name for f in dataclasses.fields(arr)] == [ + "metadata", + "store_path", + "codec_pipeline", + "_chunk_grid", + "config", + ] + + +@pytest.mark.parametrize( + "context", + [ + None, + Context(codec_pipeline=BatchedCodecPipeline), + Context(data_types=data_type_registry), + ], +) +def test_v3_from_dict_reports_errors_in_field_order(context: Context | None) -> None: + """Unless a context limits codecs, codecs are parsed after the other fields, + so a document with several bad fields reports the chunk grid first, as it + always has.""" + doc: dict[str, Any] = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "uint8", + "chunk_grid": {"name": "nope", "configuration": {}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "nope"}], + } + with pytest.raises(ValueError, match="chunk grid"): + ArrayV3Metadata.from_dict(doc, context=context) + + +def test_a_context_rechecks_only_the_kinds_it_limits() -> None: + """A context that limits one kind leaves the other as given, entering it + (constructor, with_metadata) and leaving it (with_context), for an instance + no registry holds.""" + codecs_only = Context(codecs={"bytes": BytesCodec}) + declared = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type=_UnregisteredUInt16(), + chunk_grid=(4,), + codecs=[BytesCodec()], + ) + legacy = declared.metadata_legacy + bindings = [ + AsyncArrayIO(legacy, StorePath(MemoryStore()), context=codecs_only), + declared.with_context(codecs_only).with_metadata(legacy), + declared.with_context(codecs_only).with_context(None), + ] + for binding in bindings: + assert type(binding.metadata_legacy.dtype) is _UnregisteredUInt16 + + registry = DataTypeRegistry() + registry.register(UInt16._zarr_v3_name, UInt16) + data_types_only = Context(data_types=registry) + unregistered_codec = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint16", + chunk_grid=(4,), + codecs=[_UnregisteredBytes()], + ) + for binding in ( + unregistered_codec.with_context(data_types_only), + unregistered_codec.with_context(data_types_only).with_context(None), + ): + assert type(binding.metadata_legacy.codecs[0]) is _UnregisteredBytes # type: ignore[union-attr] + + +@dataclass(frozen=True) +class _LocalBytes(BytesCodec): + """A serializer only a context holds, under a name of its own.""" + + def to_dict(self) -> dict[str, JSON]: + return {"name": "local_bytes", "configuration": {"endian": "little"}} + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + return cls(endian="little") + + +async def test_repair_before_a_write_reads_under_the_context() -> None: + """The repair check before a chunk write reads the stored document under the + binding's context, so a codec only the context holds does not stop a write + to an array whose stored chunk size needs a repair.""" + store = MemoryStore() + doc: dict[str, Any] = { + "zarr_format": 3, + "node_type": "array", + "shape": [3], + "data_type": "uint16", + "fill_value": 0, + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [0]}}, + "chunk_key_encoding": {"name": "default"}, + "codecs": [{"name": "local_bytes", "configuration": {"endian": "little"}}], + "attributes": {}, + } + prototype = default_buffer_prototype() + await store.set("a/zarr.json", prototype.buffer.from_bytes(json.dumps(doc).encode())) + context = Context(codecs={"local_bytes": _LocalBytes}) + with pytest.warns(ZarrUserWarning, match="is read as"): + arr = AsyncArray(cast("Any", doc), StorePath(store, "a"), context=context) + await arr.setitem(slice(None), np.array([7, 8, 9], dtype="uint16")) + np.testing.assert_array_equal(await arr.getitem(slice(None)), [7, 8, 9]) + + +def test_closed_codec_scope_rejects_an_instance_of_another_class() -> None: + """An instance of another class than the scope maps its name to is refused, not + rebuilt (that would drop state its JSON does not hold).""" + with pytest.raises(UnknownCodecError, match="maps 'bytes' to BytesCodec"): + AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[_UnregisteredBytes()], + context=Context(codecs={"bytes": BytesCodec}), + ) + + +def test_a_context_copies_a_registry_it_is_given() -> None: + registry = _scoped_uint16_registry() + context = Context(data_types=registry) + registry.register(Int16._zarr_v3_name, Int16) + assert context.data_types is not None + assert Int16._zarr_v3_name not in context.data_types.contents + assert Context(data_types=data_type_registry).data_types is data_type_registry + + +def test_a_pickled_context_keeps_the_default_registry() -> None: + """A context over the default data type registry unpickles over the + receiving process's registry, not a snapshot; another registry is copied.""" + restored = pickle.loads(pickle.dumps(Context(data_types=data_type_registry))) + assert restored.data_types is data_type_registry + own = _scoped_uint16_registry() + copied = pickle.loads(pickle.dumps(Context(data_types=own))) + assert copied.data_types is not own + assert copied.data_types is not None + assert copied.data_types.get(UInt16._zarr_v3_name) is _ScopedUInt16 + + class _TaggedAsyncArray(AsyncArray[Any]): """An `AsyncArray` subclass with an instance attribute of its own.""" tag: str - def __init__(self, metadata: Any, store_path: StorePath, tag: str = "default") -> None: - super().__init__(metadata, store_path) + def __init__( + self, metadata: Any, store_path: StorePath, config: Any = None, tag: str = "default" + ) -> None: + super().__init__(metadata, store_path, config) object.__setattr__(self, "tag", tag) @@ -780,6 +1449,13 @@ def test_async_array_subclass_state_survives_pickle_and_copy() -> None: assert restored == tagged assert set(vars(restored)) == {"_io", "tag"} + # Rebuilt through the subclass's __init__ (which takes no context), as before. + order_f = parse_array_config({"order": "F"}) + for rebuilt in (tagged.with_config(order_f), dataclasses.replace(tagged, config=order_f)): + assert type(rebuilt) is _TaggedAsyncArray + assert rebuilt.tag == "default" + assert rebuilt.config.order == "F" + async def test_copying_an_array_shares_its_metadata() -> None: """A shallow copy shares the metadata object, also for an array opened from a @@ -986,6 +1662,7 @@ def test_an_array_keeps_its_binding_class() -> None: for kept in (array, pickle.loads(pickle.dumps(array))): sync(kept.resize((6,))) assert type(kept.io) is _LoggingBinding + assert type(array.with_config({"order": "F"}).io) is _LoggingBinding @dataclass(frozen=True) @@ -1037,6 +1714,7 @@ async def test_every_route_rebinds_through_bind() -> None: built.with_metadata(built.metadata), built.with_store(MemoryStore()), built.with_config({"order": "F"}), + built.with_context(Context(codec_pipeline=_MarkedPipeline)), pickle.loads(pickle.dumps(built)), array.io, array.with_config({"order": "F"}).io, @@ -1432,30 +2110,994 @@ def test_with_config_prefers_a_binding_pipeline_over_an_assigned_one() -> None: assert isinstance(assigned.with_config({"order": "F"}).codec_pipeline, _AssignedPipeline) -class _ResolvedFirstElsewhere: - """Stands in for `parse_array_metadata` while another thread stores its own - resolution of `binding` first.""" +def test_bindings_expose_their_context() -> None: + context = Context(codec_pipeline=_MarkedPipeline) + aio = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,), context=context + ) + sync_io = ArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,), context=context + ) + assert aio.context == sync_io.context == context + assert aio.with_context(None).context is None + assert aio.with_context(Context()).context is None - def __init__(self, binding: AsyncArrayIO, elsewhere: ArrayMetadata) -> None: - self.binding = binding - self.elsewhere = elsewhere - def __call__(self, *args: Any, **kwargs: Any) -> ArrayMetadata: - resolved = parse_array_metadata(*args, **kwargs) - self.binding._resolved_cache = self.elsewhere - return resolved +@dataclass(frozen=True, init=False) +class _DecoratedAsyncArray(AsyncArray[Any]): + """A `@dataclass` subclass with a field and an `__init__` that takes no context.""" + tag: str = "default" -async def test_bindings_resolving_at_once_share_the_first_resolution() -> None: - """A binding over a model keeps the legacy metadata resolved first, so threads - resolving it at once share one object, whose attributes an array changes in place.""" - store = MemoryStore() - built = AsyncArrayIO.build_v3(store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,)) - await built.store_metadata() - binding = await AsyncArrayIO.open(store, name="a") - elsewhere = parse_array_metadata(built.metadata_legacy.to_dict(), "a") - with mock.patch( - "zarr.core.array.parse_array_metadata", _ResolvedFirstElsewhere(binding, elsewhere) - ): - assert binding.metadata_legacy is elsewhere - assert binding.metadata_legacy is elsewhere + def __init__( + self, metadata: Any, store_path: StorePath, config: Any = None, tag: str = "default" + ) -> None: + super().__init__(metadata, store_path, config) + object.__setattr__(self, "tag", tag) + + +def test_a_subclass_init_without_context_is_rebuilt() -> None: + """`dataclasses.replace` and `with_config` call a subclass `__init__` that takes + no context, decorated or not, as they always did; a decorated subclass's own + field is passed by `replace`, and `with_config` keeps a context.""" + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + order_f = parse_array_config({"order": "F"}) + context = Context(codec_pipeline=_MarkedPipeline) + for cls, replaced_tag in ((_TaggedAsyncArray, "default"), (_DecoratedAsyncArray, "mine")): + array = cls(created.metadata, created.store_path, tag="mine") + for rebuilt, tag in ( + (dataclasses.replace(array, config=order_f), replaced_tag), + (array.with_config(order_f), "default"), + ): + assert type(rebuilt) is cls + assert rebuilt.tag == tag + assert rebuilt.config.order == "F" + scoped = cls._from_io(AsyncArrayIO(created.metadata, created.store_path, context=context)) + object.__setattr__(scoped, "tag", "mine") + configured = scoped.with_config(order_f) + assert type(configured) is cls + assert configured.tag == "default" + assert configured.io.context == context + assert isinstance(configured.codec_pipeline, _MarkedPipeline) + + +def test_async_array_takes_context_by_keyword_only() -> None: + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + with pytest.raises(TypeError, match="positional"): + AsyncArray(created.metadata, created.store_path, None, Context()) # type: ignore[call-overload] + + +def test_async_array_rejects_a_model_under_any_context() -> None: + """`AsyncArray` takes legacy metadata or a mapping; a model is refused under a + context that limits names as without one.""" + model = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,) + ).metadata + closed = Context(codecs={"bytes": BytesCodec}) + for context in (None, closed): + with pytest.raises(TypeError, match="got ZarrV3ArrayMetadata"): + AsyncArray(cast("Any", model), StorePath(MemoryStore()), context=context) + + +def test_a_context_with_only_unset_kinds_binds_as_none() -> None: + aio = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,), context=Context() + ) + assert aio.context is None + + +class _BuiltBinding(AsyncArrayIO): + """A binding subclass whose own `__init__` sets up state, as a backend's might.""" + + ready: bool + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.ready = True + + +def test_builders_run_a_binding_subclass_init() -> None: + context = Context(codec_pipeline=_MarkedPipeline) + built = ( + _BuiltBinding.build_v3(MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,)), + _BuiltBinding.build_v2( + MemoryStore(), shape=(4,), dtype="uint8", chunks=(4,), context=context + ), + ) + for binding in built: + assert binding.ready + + +async def test_with_config_under_a_context_shares_metadata_and_builds_one_pipeline() -> None: + """Under a context, `with_config` keeps the context and its pipeline, builds no + other pipeline, and shares the metadata object, as it does with no context.""" + closed = Context(codecs={"bytes": BytesCodec}, codec_pipeline=_MarkedPipeline) + for context in (None, closed): + array = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[BYTES_JSON], + context=context, + ).to_array() + configured = array.with_config({"order": "F"}) + assert configured.metadata is array.metadata + assert configured.io.context == context + array = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[BYTES_JSON], + context=closed, + ).to_array() + with zarr.config.set({"codec_pipeline.path": "no.such.Pipeline"}): + configured = array.with_config({"order": "F"}) + assert isinstance(configured.codec_pipeline, _MarkedPipeline) + + +def test_sync_bindings_take_a_context(tmp_path: Path) -> None: + context = Context(codec_pipeline=_MarkedPipeline) + built = ArrayIO.build_v2( + LocalStore(tmp_path), name="a", shape=(4,), dtype="uint8", chunks=(4,), context=context + ) + built.store_metadata() + opened = ArrayIO.open(LocalStore(tmp_path), name="a", context=context) + assert built.context == opened.context == context + assert opened.with_context(None).context is None + assert isinstance(opened.to_array().async_array.codec_pipeline, _MarkedPipeline) + + +@pytest.mark.parametrize("context", [None, Context(data_types=data_type_registry)]) +@pytest.mark.parametrize( + "native", [np.dtype(">u2"), np.dtype([("a", ">i4"), ("n", [("c", ">u2")])])] +) +def test_a_data_type_scope_keeps_the_byte_order( + context: Context | None, native: np.dtype[Any] +) -> None: + """Zarr format 3 data type JSON has no byte order; entering a data type scope, or + leaving it, keeps a big-endian data type (or field) big-endian, as no context does.""" + built = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type=native, chunk_grid=(4,), context=context + ) + created = zarr.create_array(MemoryStore(), shape=(4,), dtype=native) + bound = AsyncArrayIO(created.metadata, created.store_path, context=context) + for binding in (built, bound, bound.with_context(None)): + assert binding.metadata_legacy.dtype.to_native_dtype() == native + + +@dataclass(frozen=True) +class _TunedZstd(ZstdCodec): + """A zstd codec with state its JSON does not hold.""" + + threads: int = 1 + + +def test_a_closed_codec_scope_keeps_admitted_instances_whole() -> None: + """An instance of the class the scope maps its name to is bound as given on + every route, state its JSON does not hold included, also nested in sharding.""" + context = Context(codecs={"bytes": BytesCodec, "zstd": _TunedZstd}) + tuned = _TunedZstd(level=3, threads=8) + built = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[BytesCodec(), tuned], + context=context, + ) + plain = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,), codecs=[BytesCodec(), tuned] + ) + legacy = plain.metadata_legacy + for binding in ( + built, + plain.with_context(context), + AsyncArrayIO(legacy, plain.store_path, context=context), + AsyncArray(legacy, plain.store_path, context=context).io, + ): + assert binding.metadata_legacy.codecs[1] == tuned # type: ignore[union-attr] + assert binding.metadata_legacy.codecs[1].threads == 8 # type: ignore[union-attr] + sharded = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[ShardingCodec(chunk_shape=(2,), codecs=[BytesCodec(), tuned])], + ) + scoped = sharded.with_context(Context(codecs={"sharding_indexed": ShardingCodec})) + assert scoped.metadata_legacy.codecs[0].codecs[1].threads == 8 # type: ignore[union-attr] + + +@dataclass(frozen=True, kw_only=True) +class _AliasRegistry(DataTypeRegistry): + """A registry that matches a name of its own.""" + + def match_json(self, data: Any, *, zarr_format: ZarrFormat) -> Any: + return super().match_json("float32" if data == "real32" else data, zarr_format=zarr_format) + + +def test_a_context_keeps_the_class_of_a_registry_it_copies() -> None: + registry = _AliasRegistry() + registry.register(Float32._zarr_v3_name, Float32) + context = Context(data_types=registry) + assert type(context.data_types) is _AliasRegistry + assert type(parse_dtype("real32", zarr_format=3, context=context)) is Float32 + + +def test_a_closed_data_type_scope_rejects_an_instance_of_another_class() -> None: + """A data type instance is admitted only if the scope reads its JSON as its own + class, as a codec instance is.""" + context = Context(data_types=_scoped_uint16_registry()) + with pytest.raises(ValueError, match="reads its JSON as, _ScopedUInt16"): + parse_dtype(UInt16(), zarr_format=3, context=context) + assert type(parse_dtype(_ScopedUInt16(), zarr_format=3, context=context)) is _ScopedUInt16 + + +def test_matching_in_a_registry_does_not_copy_it(monkeypatch: pytest.MonkeyPatch) -> None: + import zarr.core.context as context_module + + registry = _scoped_uint16_registry() + monkeypatch.setattr( + context_module, + "_registry_snapshot", + mock.Mock(side_effect=AssertionError("the registry was copied")), + ) + assert type(registry.match_json("uint16", zarr_format=3)) is _ScopedUInt16 + + +def test_with_config_chooses_the_pipeline_anew_unless_the_context_chose_it() -> None: + """Under a context that sets no pipeline, `with_config` follows the global + configuration, as with no context; a context's own pipeline is kept.""" + fused = "zarr.core.codec_pipeline.FusedCodecPipeline" + for context, expected in ( + (Context(data_types=_scoped_uint16_registry()), FusedCodecPipeline), + (Context(codec_pipeline=_MarkedPipeline), _MarkedPipeline), + ): + array = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint16", chunk_grid=(4,), context=context + ).to_array() + with zarr.config.set({"codec_pipeline.path": fused}): + configured = array.with_config({"order": "F"}) + assert type(configured.codec_pipeline) is expected + + +@dataclass(frozen=True) +class _LazyEntryPoint: + """An entry point that loads a data type class.""" + + loaded: type[Any] + + def load(self) -> type[Any]: + return self.loaded + + +def test_a_context_compares_the_same_before_and_after_its_registry_loads() -> None: + registry = DataTypeRegistry() + registry._lazy_load_list.append(cast("Any", _LazyEntryPoint(_ScopedUInt16))) + context = Context(data_types=registry) + binding = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint16", chunk_grid=(4,), context=context + ) + _ = binding.metadata_legacy + assert context == Context(data_types=registry) + + +def test_a_closed_data_type_scope_checks_struct_field_classes() -> None: + """A structured data type instance is admitted only if the scope reads each of + its fields as that field's own class.""" + context = Context(data_types=_scoped_uint16_registry_with_struct()) + with pytest.raises(ValueError, match="reads its JSON as"): + parse_dtype(Struct(fields=(("a", UInt16()),)), zarr_format=3, context=context) + admitted = Struct(fields=(("a", _ScopedUInt16()),)) + assert parse_dtype(admitted, zarr_format=3, context=context) is admitted + + +def _scoped_uint16_registry_with_struct() -> DataTypeRegistry: + registry = _scoped_uint16_registry() + registry.register(Struct._zarr_v3_name, Struct) + return registry + + +def test_leaving_a_codec_scope_keeps_default_codecs_whole() -> None: + """Leaving a scope rebinds only codecs the scope bound to other classes: a + default-class codec is kept whole, with the codecs nested in it.""" + tuned = _TunedZstd(level=3, threads=8) + sharded = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[ShardingCodec(chunk_shape=(2,), codecs=[BytesCodec(), tuned])], + ) + left = sharded.with_context(Context(codecs={"sharding_indexed": ShardingCodec})).with_context( + None + ) + assert left.metadata_legacy.codecs[0].codecs[1].threads == 8 # type: ignore[union-attr] + + +def test_rebinding_resolved_metadata_reads_other_classes_by_name() -> None: + """Metadata already resolved, entering a scope that maps a name to another class + than its codec's, is read again as the scope's class (a builder refuses such an + instance).""" + plain = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[BytesCodec(), _TunedZstd(level=3, threads=8)], + ) + scoped = plain.with_context(Context(codecs={"bytes": BytesCodec, "zstd": ZstdCodec})) + assert type(scoped.metadata_legacy.codecs[1]) is ZstdCodec # type: ignore[union-attr] + + +class _ScopedShuffle(Shuffle): + """A shuffle codec class of a scope's own.""" + + +async def test_a_scoped_numcodecs_class_survives_fitting_to_the_array() -> None: + """A numcodecs codec fitted to the array's data type stays the scope's class.""" + store = MemoryStore() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "int32", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "numcodecs.shuffle", "configuration": {}}, + ], + "attributes": {}, + } + await store.set( + "a/zarr.json", default_buffer_prototype().buffer.from_bytes(json.dumps(document).encode()) + ) + context = Context(codecs={"bytes": BytesCodec, "numcodecs.shuffle": _ScopedShuffle}) + binding = await AsyncArrayIO.open(store, name="a", context=context) + assert type(binding.metadata_legacy.codecs[1]) is _ScopedShuffle # type: ignore[union-attr] + + +async def test_a_scoped_write_passes_a_recorded_document_its_scope_cannot_read() -> None: + """A document the store's cache records as needing a repair, which the writing + binding's scope cannot read, describes another array: the write goes ahead.""" + store = MemoryStore() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [3], + "data_type": "int32", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [0]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + } + await store.set( + "a/zarr.json", default_buffer_prototype().buffer.from_bytes(json.dumps(document).encode()) + ) + with pytest.warns(ZarrUserWarning): + await open_array(store=store, path="a") + context = Context(data_types=_scoped_uint16_registry()) + binding = AsyncArrayIO.build_v3( + store, name="a", shape=(3,), data_type="uint16", chunk_grid=(3,), context=context + ) + await binding.setitem(slice(None), np.array([1, 2, 3], dtype="uint16")) + # Another array's layout: its document is left as it was. + stored = await store.get("a/zarr.json", prototype=default_buffer_prototype()) + assert stored is not None + assert json.loads(stored.to_bytes())["chunk_grid"]["configuration"]["chunk_shape"] == [0] + + +class _ResubmittingAsyncArray(AsyncArray[Any]): + """A subclass whose `__init__` gives its metadata as a dict.""" + + def __init__(self, metadata: Any, store_path: StorePath, config: Any = None) -> None: + super().__init__(metadata.to_dict(), store_path, config) + + +def test_with_config_resolves_substituted_metadata_under_the_context() -> None: + context = Context(data_types=_scoped_uint16_registry()) + binding = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint16", chunk_grid=(4,), context=context + ) + array = _ResubmittingAsyncArray._from_io(binding) + configured = array.with_config({"order": "F"}) + assert type(configured.metadata.dtype) is _ScopedUInt16 + assert configured.io.context == context + + +def test_bindings_take_context_by_keyword_only() -> None: + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + for binding_class in (AsyncArrayIO, ArrayIO): + with pytest.raises(TypeError, match="positional"): + binding_class(created.metadata, created.store_path, None, Context()) # type: ignore[call-arg] + + +def test_with_config_without_a_context_runs_the_constructor() -> None: + """Without a context, `with_config` builds the array through its constructor, as + it always has.""" + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + array = AsyncArray(created.metadata, created.store_path) + with mock.patch.object( + AsyncArray, "__init__", autospec=True, side_effect=AsyncArray.__init__ + ) as init: + array.with_config({"order": "F"}) + assert init.call_count == 1 + + +def test_the_builder_keeps_a_scoped_numcodecs_class_through_fitting() -> None: + """`build_array_metadata_v3` under a codec scope returns the scope's class for a + codec that fitting to the array rebuilt, as reading a document does.""" + context = Context(codecs={"bytes": BytesCodec, "numcodecs.shuffle": _ScopedShuffle}) + metadata = build_array_metadata_v3( + shape=(4,), + data_type="int32", + chunk_grid=(4,), + codecs=[BYTES_JSON, {"name": "numcodecs.shuffle", "configuration": {}}], + context=context, + ) + assert type(metadata.codecs[1]) is _ScopedShuffle + + +def test_a_name_limiting_scope_shares_metadata_it_does_not_change() -> None: + """Under a scope that reads every name as the metadata already holds it, an + `AsyncArray` binds the object it is given, as with no context.""" + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(4,)) + context = Context(data_types=data_type_registry) + array = AsyncArray(created.metadata, created.store_path, context=context) + assert array.metadata is created.metadata + + +def test_with_config_builds_its_pipeline_over_metadata_resolved_under_the_scope() -> None: + """A subclass `__init__` that substitutes metadata read under the defaults: the + array is resolved again under its context, and its pipeline is built over that.""" + context = Context(codecs={"bytes": _UnregisteredBytes}) + binding = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[BYTES_JSON], + context=context, + ) + configured = _ResubmittingAsyncArray._from_io(binding).with_config({"order": "F"}) + pipeline = cast("BatchedCodecPipeline", configured.codec_pipeline) + assert type(pipeline.array_bytes_codec) is _UnregisteredBytes + + +class _LocalPipelineBinding(AsyncArrayIO): + """A binding whose `_bind` chooses its own pipeline for a local store only.""" + + def _bind(self, *args: Any) -> None: + *rest, pipeline_class = args + local = isinstance(rest[1].store, LocalStore) + bind: Any = super()._bind + bind(*rest, _OwnPipeline if local else pipeline_class) + + +def test_a_binding_made_from_another_gets_its_contexts_pipeline(tmp_path: Path) -> None: + """Under a context that chose a pipeline, a binding made from one whose `_bind` + chose another class for its store gets the context's class again.""" + context = Context(codec_pipeline=_MarkedPipeline) + local = _LocalPipelineBinding.build_v3( + LocalStore(tmp_path), shape=(4,), data_type="uint8", chunk_grid=(4,), context=context + ) + assert local._pipeline_class is _OwnPipeline + assert local.with_store(MemoryStore())._pipeline_class is _MarkedPipeline + + +class _ShapePipelineBinding(AsyncArrayIO): + """A binding whose `_bind` chooses its own pipeline for small arrays only.""" + + def _bind(self, *args: Any) -> None: + source, *rest, pipeline_class = args + bind: Any = super()._bind + bind(source, *rest, _OwnPipeline if source.shape[0] <= 4 else pipeline_class) + + +async def test_a_resize_runs_the_pipeline_its_binding_records() -> None: + """A resize after which `_bind` chooses another pipeline class runs that class.""" + context = Context(codec_pipeline=_MarkedPipeline) + array = _ShapePipelineBinding.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,), context=context + ).to_array() + assert isinstance(array.codec_pipeline, _OwnPipeline) + await array.resize((8,)) + assert array.io._pipeline_class is _MarkedPipeline + assert isinstance(array.codec_pipeline, _MarkedPipeline) + + +async def test_a_scoped_write_repairs_a_matching_document_its_scope_cannot_read() -> None: + """A recorded document the scope cannot read is compared by its layout, read + under the defaults, as with no context: one laid out as the binding is gets + its repair stored before the write.""" + store = MemoryStore() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [3], + "data_type": "int32", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [0]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + } + await store.set( + "a/zarr.json", default_buffer_prototype().buffer.from_bytes(json.dumps(document).encode()) + ) + with pytest.warns(ZarrUserWarning): + await open_array(store=store, path="a") + context = Context(data_types=_scoped_uint16_registry()) + binding = AsyncArrayIO.build_v3( + store, name="a", shape=(3,), data_type="uint16", chunk_grid=(1,), context=context + ) + await binding.setitem(slice(None), np.array([1, 2, 3], dtype="uint16")) + stored = await store.get("a/zarr.json", prototype=default_buffer_prototype()) + assert stored is not None + chunk_grid = json.loads(stored.to_bytes())["chunk_grid"] + assert chunk_grid["configuration"]["chunk_shape"] == [1] + + +def test_arrays_expose_their_context() -> None: + context = Context(codec_pipeline=_MarkedPipeline) + binding = ArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,), context=context + ) + array = binding.to_array() + assert array.async_array.context == array.io.context == context + + +class _OrderPipelineBinding(AsyncArrayIO): + """A binding whose `_bind` chooses its own pipeline for C order only.""" + + def _bind(self, source: Any, store_path: Any, config: Any, *rest: Any) -> None: + *context, pipeline_class = rest + chosen = _OwnPipeline if config.order == "C" else pipeline_class + bind: Any = super()._bind + bind(source, store_path, config, *context, chosen) + + +def test_setting_a_config_runs_the_pipeline_its_binding_records() -> None: + """Setting an array's config, after which `_bind` chooses another pipeline class, + runs that class.""" + context = Context(codec_pipeline=_MarkedPipeline) + array = _OrderPipelineBinding.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,), context=context + ).to_array() + assert isinstance(array.codec_pipeline, _OwnPipeline) + object.__setattr__(array, "config", parse_array_config({"order": "F"})) + assert array.io._pipeline_class is _MarkedPipeline + assert isinstance(array.codec_pipeline, _MarkedPipeline) + + +async def test_a_scoped_write_reads_a_newer_document_its_scope_cannot_read() -> None: + """The document stored now is laid out under the defaults when the binding's + scope cannot read it, as the recorded one is: a newer valid document another + writer stored lets the write go ahead, as with no context.""" + store = MemoryStore() + prototype = default_buffer_prototype() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [3], + "data_type": "uint16", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [0]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + } + await store.set("a/zarr.json", prototype.buffer.from_bytes(json.dumps(document).encode())) + with pytest.warns(ZarrUserWarning): + await open_array(store=store, path="a") + newer = { + **document, + "data_type": "int32", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [1]}}, + } + await store.set("a/zarr.json", prototype.buffer.from_bytes(json.dumps(newer).encode())) + context = Context(data_types=_scoped_uint16_registry()) + binding = AsyncArrayIO.build_v3( + store, name="a", shape=(3,), data_type="uint16", chunk_grid=(1,), context=context + ) + await binding.setitem(slice(None), np.array([1, 2, 3], dtype="uint16")) + assert {"a/c/0", "a/c/1", "a/c/2"} <= await _store_keys(store) + stored = await store.get("a/zarr.json", prototype=default_buffer_prototype()) + assert stored is not None + assert json.loads(stored.to_bytes()) == newer + + +async def test_a_scoped_write_over_a_document_no_registry_reads_raises() -> None: + """A recorded document that neither the writing binding's scope nor the + defaults can read refuses the write, storing nothing.""" + store = MemoryStore() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [3], + "data_type": "uint16", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [0]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "local_bytes", "configuration": {"endian": "little"}}], + "attributes": {}, + } + with pytest.warns(ZarrUserWarning): + AsyncArray( + cast("Any", document), + StorePath(store, "a"), + context=Context(codecs={"local_bytes": _LocalBytes}), + ) + binding = AsyncArrayIO.build_v3( + store, + name="a", + shape=(3,), + data_type="uint16", + chunk_grid=(1,), + context=Context(data_types=_scoped_uint16_registry()), + ) + with pytest.raises(UnknownCodecError): + await binding.setitem(slice(None), np.array([1, 2, 3], dtype="uint16")) + assert await _store_keys(store) == set() + + +def test_without_a_context_a_derived_binding_runs_the_class_bind_chose(tmp_path: Path) -> None: + """With no context, the pipeline class a `_bind` chose is the one bindings made + from that binding are given.""" + local = _LocalPipelineBinding.build_v3( + LocalStore(tmp_path), shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + assert local._pipeline_class is _OwnPipeline + assert local.with_store(MemoryStore())._pipeline_class is _OwnPipeline + + +def test_leaving_a_scope_fails_for_a_name_the_defaults_do_not_know() -> None: + binding = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint16", + chunk_grid=(4,), + codecs=[{"name": "local_bytes", "configuration": {"endian": "little"}}], + context=Context(codecs={"local_bytes": _LocalBytes}), + ) + with pytest.raises(UnknownCodecError, match="local_bytes"): + binding.with_context(None) + + +@dataclass(frozen=True) +class _PlainZstd(ZstdCodec): + """A zstd class whose constructor takes no `checksum`.""" + + def __init__(self, *, level: int = 0) -> None: + super().__init__(level=level) + + +async def test_a_scoped_write_reads_a_document_a_scope_class_refuses_under_the_defaults() -> None: + """A document a scope's class refuses to read (its constructor rejects the stored + configuration) is laid out under the defaults, as with no context.""" + store = MemoryStore() + prototype = default_buffer_prototype() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [3], + "data_type": "uint16", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [0]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "zstd", "configuration": {"level": 0, "checksum": True}}, + ], + } + await store.set("a/zarr.json", prototype.buffer.from_bytes(json.dumps(document).encode())) + with pytest.warns(ZarrUserWarning): + await open_array(store=store, path="a") + context = Context(codecs={"bytes": BytesCodec, "zstd": _PlainZstd}) + binding = AsyncArrayIO.build_v3( + store, + name="a", + shape=(3,), + data_type="uint16", + chunk_grid=(1,), + codecs=[BYTES_JSON], + context=context, + ) + await binding.setitem(slice(None), np.array([1, 2, 3], dtype="uint16")) + assert {"a/c/0", "a/c/1", "a/c/2"} <= await _store_keys(store) + + +@dataclass(frozen=True) +class _ScopedGzip(GzipCodec): + """A gzip class only a scope knows, under a name of its own.""" + + def to_dict(self) -> dict[str, JSON]: + return {"name": "scoped_gzip", "configuration": {"level": self.level}} + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + configuration = cast("dict[str, JSON]", data["configuration"]) + return cls(level=cast("int", configuration["level"])) + + +async def test_a_scoped_write_reads_each_document_under_its_scope_first() -> None: + """The current document is read under the binding's scope even when the recorded + one needed the defaults.""" + store = MemoryStore() + prototype = default_buffer_prototype() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [3], + "data_type": "uint16", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [0]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "gzip", "configuration": {"level": 1}}, + ], + } + await store.set("a/zarr.json", prototype.buffer.from_bytes(json.dumps(document).encode())) + with pytest.warns(ZarrUserWarning): + await open_array(store=store, path="a") + newer = { + **document, + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [1]}}, + "codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "scoped_gzip", "configuration": {"level": 1}}, + ], + } + await store.set("a/zarr.json", prototype.buffer.from_bytes(json.dumps(newer).encode())) + context = Context(codecs={"bytes": BytesCodec, "scoped_gzip": _ScopedGzip}) + binding = AsyncArrayIO.build_v3( + store, + name="a", + shape=(3,), + data_type="uint16", + chunk_grid=(1,), + codecs=[BYTES_JSON], + context=context, + ) + await binding.setitem(slice(None), np.array([1, 2, 3], dtype="uint16")) + assert {"a/c/0", "a/c/1", "a/c/2"} <= await _store_keys(store) + + +def test_a_codec_scope_held_under_an_alias_admits_no_instance() -> None: + """A scope names codecs by the name each writes: one that holds `BytesCodec` + only as `endian` admits no `BytesCodec` instance (it writes `bytes`, which the + scope could not read back), whether built or entering the scope.""" + context = Context(codecs={"endian": BytesCodec}) + with pytest.raises(UnknownCodecError, match="'bytes' is not in this context"): + AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint16", + chunk_grid=(4,), + codecs=[BytesCodec()], + context=context, + ) + with pytest.raises(UnknownCodecError, match="'bytes' is not in this context") as raised: + AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint16", chunk_grid=(4,), codecs=[BytesCodec()] + ).with_context(context) + assert raised.value.__notes__ == ["'bytes' is the name a BytesCodec writes."] + + +async def test_a_rebind_that_switches_pipeline_builds_it_on_first_use() -> None: + """A resize has stored its shape before the binding rebinds, so a pipeline class + `_bind` switches to is built when next used, never failing the resize.""" + context = Context(codec_pipeline=_MarkedPipeline) + array = _ShapePipelineBinding.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,), context=context + ).to_array() + await array.resize((8,)) + assert array.io._codec_pipeline_cache is None + assert isinstance(array.codec_pipeline, _MarkedPipeline) + + +class _NamedBinding(AsyncArrayIO): + """A binding whose constructor takes a signature of its own.""" + + def __init__(self, name: str) -> None: + raise AssertionError("derivations do not run the constructor") + + +def test_with_context_derives_without_the_constructor() -> None: + """`with_context` makes its binding as the other derivations do, so a subclass + whose constructor has its own signature derives, and a legacy source is + never shared with the original binding.""" + built = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type="uint16", chunk_grid=(4,), codecs=[BytesCodec()] + ) + binding = object.__new__(_NamedBinding) + binding._bind( + built._source, + built.store_path, + built.config, + None, + built._pipeline_class, + ) + context = Context(codecs={"bytes": BytesCodec}) + derived = binding.with_context(context) + assert type(derived) is _NamedBinding + assert derived.context == context + assert derived.metadata_legacy == binding.metadata_legacy + assert derived._source is not binding._source + assert type(derived.with_context(None)) is _NamedBinding + + +class _ResolvedFirstElsewhere: + """Stands in for `parse_array_metadata` while another thread stores its own + resolution of `binding` first.""" + + def __init__(self, binding: AsyncArrayIO, elsewhere: ArrayMetadata) -> None: + self.binding = binding + self.elsewhere = elsewhere + + def __call__(self, *args: Any, **kwargs: Any) -> ArrayMetadata: + resolved = parse_array_metadata(*args, **kwargs) + self.binding._resolved_cache = self.elsewhere + return resolved + + +async def test_bindings_resolving_at_once_share_the_first_resolution() -> None: + """A binding over a model keeps the legacy metadata resolved first, so threads + resolving it at once share one object, whose attributes an array changes in place.""" + store = MemoryStore() + built = AsyncArrayIO.build_v3(store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,)) + await built.store_metadata() + binding = await AsyncArrayIO.open(store, name="a") + elsewhere = parse_array_metadata(built.metadata_legacy.to_dict(), "a") + with mock.patch( + "zarr.core.array.parse_array_metadata", _ResolvedFirstElsewhere(binding, elsewhere) + ): + assert binding.metadata_legacy is elsewhere + assert binding.metadata_legacy is elsewhere + + +def test_a_data_type_scope_rereads_without_unstable_specification_warnings() -> None: + """Reading a resolved data type again under a scope, or leaving one, writes + nothing, so it gives no warning of an unstable specification (warnings are + errors in this suite); only storing the metadata does, as with no context.""" + registry = DataTypeRegistry() + registry.register(NullTerminatedBytes._zarr_v3_name, NullTerminatedBytes) + context = Context(data_types=registry) + data_type = NullTerminatedBytes(length=4) + scoped = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type=data_type, chunk_grid=(4,), context=context + ) + left = scoped.with_context(None) + entered = AsyncArrayIO.build_v3( + MemoryStore(), shape=(4,), data_type=data_type, chunk_grid=(4,) + ).with_context(context) + assert scoped.metadata_legacy.dtype == left.metadata_legacy.dtype == data_type + assert entered.metadata_legacy.dtype == data_type + + +def test_bindings_that_run_different_pipelines_are_unequal() -> None: + """A binding's pipeline class is part of it: bindings over one document made + under different global pipeline settings run different pipelines, so they are + unequal.""" + store = MemoryStore() + default = AsyncArrayIO.build_v3(store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,)) + with zarr.config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + fused = AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + assert default._pipeline_class is not fused._pipeline_class + assert default != fused + + +def test_a_data_type_scope_leaves_the_warning_filters_alone() -> None: + """Silencing the unstable-specification warning for a reread touches no global + warning state: a context-free warning shown once stays shown once around + scoped reads, and the filters are unchanged.""" + registry = DataTypeRegistry() + registry.register(NullTerminatedBytes._zarr_v3_name, NullTerminatedBytes) + context = Context(data_types=registry) + data_type = NullTerminatedBytes(length=4) + with warnings.catch_warnings(record=True) as shown: + warnings.simplefilter("default") + filters = list(warnings.filters) + for _ in range(3): + data_type.to_json(zarr_format=3) + parse_dtype(data_type, zarr_format=3, context=context) + assert warnings.filters == filters + assert [type(w.message) for w in shown] == [UnstableSpecificationWarning] + + +class _FixedSignatureBinding(AsyncArrayIO): + """A binding subclass whose constructor was written without a `context`.""" + + def __init__( + self, metadata: Any, store_path: StorePath, config: ArrayConfigLike | None = None + ) -> None: + super().__init__(metadata, store_path, config) + + +async def test_builders_and_open_need_no_context_keyword_without_a_context() -> None: + store = MemoryStore() + built = _FixedSignatureBinding.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + await built.store_metadata() + v2 = _FixedSignatureBinding.build_v2(MemoryStore(), shape=(4,), dtype="uint8", chunks=(4,)) + opened = await _FixedSignatureBinding.open(store, name="a") + assert type(built) is type(v2) is type(opened) is _FixedSignatureBinding + assert opened.metadata_legacy == built.metadata_legacy + + +async def test_an_opened_binding_entering_a_codec_scope_resolves_its_classes() -> None: + """A binding over a model (one `open` made) entering a codec scope resolves its + stored names through the scope when its legacy metadata is first built.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,), codecs=[BytesCodec()] + ).store_metadata() + opened = await AsyncArrayIO.open(store, name="a") + scoped = opened.with_context(Context(codecs={"bytes": _UnregisteredBytes})) + assert [type(codec) for codec in scoped.metadata_legacy.codecs] == [_UnregisteredBytes] # type: ignore[union-attr] + + +async def test_an_opened_binding_entering_a_scope_without_its_codec_fails_on_use() -> None: + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + opened = await AsyncArrayIO.open(store, name="a") + scoped = opened.with_context(Context(codecs={"zstd": ZstdCodec})) + with pytest.raises(UnknownCodecError, match="'bytes' is not in this context"): + await scoped.getitem(slice(None)) + + +def test_a_scope_class_refusing_a_codec_configuration_names_it() -> None: + built = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[BytesCodec(), ZstdCodec(level=0, checksum=True)], + ) + context = Context(codecs={"bytes": BytesCodec, "zstd": _PlainZstd}) + with pytest.raises(ZarrMetadataValidationError, match="'zstd'.*_PlainZstd, which refuses"): + built.with_context(context) + # The same refusal read from JSON reports alike. + with pytest.raises(ZarrMetadataValidationError, match="'zstd'.*_PlainZstd, which refuses"): + AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[ + {"name": "bytes"}, + {"name": "zstd", "configuration": {"level": 0, "checksum": True}}, + ], + context=context, + ) + + +class _RenamingZstd(ZstdCodec): + """A zstd class that writes a name of its own.""" + + def to_dict(self) -> dict[str, JSON]: + return {**super().to_dict(), "name": "my.zstd"} + + +def test_a_scope_class_rebuilding_a_codec_must_read_back_in_the_scope() -> None: + built = AsyncArrayIO.build_v3( + MemoryStore(), + shape=(4,), + data_type="uint8", + chunk_grid=(4,), + codecs=[BytesCodec(), ZstdCodec()], + ) + context = Context(codecs={"bytes": BytesCodec, "zstd": _RenamingZstd}) + with pytest.raises(UnknownCodecError, match="'my.zstd' is not in this context") as raised: + built.with_context(context) + assert raised.value.__notes__ == [ + "'my.zstd' is the name a _RenamingZstd writes, read from 'zstd'." + ] diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 785f089513..6ffaa4dbfd 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -20,6 +20,7 @@ from zarr import Array from zarr.core.buffer import default_buffer_prototype from zarr.core.chunk_grids import ChunkGrid +from zarr.core.codec_pipeline import BatchedCodecPipeline, FusedCodecPipeline from zarr.core.common import ceildiv_int from zarr.core.indexing import ( BasicSelection, @@ -38,6 +39,7 @@ oindex_set, replace_ellipsis, ) +from zarr.experimental import Context from zarr.registry import get_ndbuffer_class from zarr.storage import MemoryStore, StorePath @@ -2401,10 +2403,16 @@ def test_set_selection_rejects_value_with_wrong_rank( """ a = np.zeros((4, 4), dtype=np.int32) value = np.arange(np.prod(value_shape), dtype=np.int32).reshape(value_shape) - with zarr.config.set({"codec_pipeline.path": pipeline_path}): - z = zarr_array_from_numpy_array(store, a, chunk_shape=(2, 2), shards=shards) - with pytest.raises(ValueError, match="Attempting to set|shape mismatch"): - getattr(z, kind)[selection] = value + pipeline_contexts = { + "zarr.core.codec_pipeline.BatchedCodecPipeline": BatchedCodecPipeline, + "zarr.core.codec_pipeline.FusedCodecPipeline": FusedCodecPipeline, + } + context = Context(codec_pipeline=pipeline_contexts[pipeline_path]) + z = zarr_array_from_numpy_array(store, a, chunk_shape=(2, 2), shards=shards) + # The pipeline is chosen per binding, not through the global config. + z = z.io.with_context(context).to_array() + with pytest.raises(ValueError, match="Attempting to set|shape mismatch"): + getattr(z, kind)[selection] = value def test_iter_chunk_regions(): diff --git a/tests/test_pipeline_parity.py b/tests/test_pipeline_parity.py index 81b9ce2e68..2cdefd57d3 100644 --- a/tests/test_pipeline_parity.py +++ b/tests/test_pipeline_parity.py @@ -51,13 +51,16 @@ SubchunkWriteOrder, ) from zarr.codecs.transpose import TransposeCodec -from zarr.core.config import config as zarr_config +from zarr.core.codec_pipeline import BatchedCodecPipeline, FusedCodecPipeline from zarr.errors import ZarrUserWarning +from zarr.experimental import Context from zarr.storage import MemoryStore if TYPE_CHECKING: from collections.abc import Callable, Iterator + from zarr.types import AnyArray + # --------------------------------------------------------------------------- # Reference helpers @@ -254,6 +257,23 @@ def _ignore_sharding_combo_warning() -> Iterator[None]: yield +_PIPELINE_CONTEXTS = { + "zarr.core.codec_pipeline.BatchedCodecPipeline": Context(codec_pipeline=BatchedCodecPipeline), + "zarr.core.codec_pipeline.FusedCodecPipeline": Context(codec_pipeline=FusedCodecPipeline), +} + + +def _under_pipeline(arr: AnyArray, pipeline_path: str) -> AnyArray: + """The same array binding running its IO under the chosen pipeline, + selected per binding through a context instead of the global config.""" + with _ignore_sharding_combo_warning(): + rebound = arr.io.with_context(_PIPELINE_CONTEXTS[pipeline_path]).to_array() + # The pipeline derives lazily; build it here, inside the suppression, + # so the sharding advisory does not fire from the first write. + _ = rebound.async_array.codec_pipeline + return rebound + + def _write_under_pipeline( pipeline_path: str, codec_kwargs: CodecConfig, @@ -271,26 +291,25 @@ def _write_under_pipeline( # endian-sensitive int dtype). Merge so the override wins without a dup kwarg. create_kwargs = {"dtype": "float64", **array_layout, **codec_kwargs} store = MemoryStore() - with zarr_config.set({"codec_pipeline.path": pipeline_path}): - with _ignore_sharding_combo_warning(): - arr = zarr.create_array( - store=store, - fill_value=0, - config={"write_empty_chunks": write_empty_chunks}, - **create_kwargs, - ) - for sel, val in sequence: - arr[sel] = val - contents = arr[...] + with _ignore_sharding_combo_warning(): + arr = zarr.create_array( + store=store, + fill_value=0, + config={"write_empty_chunks": write_empty_chunks}, + **create_kwargs, + ) + arr = _under_pipeline(arr, pipeline_path) + for sel, val in sequence: + arr[sel] = val + contents = arr[...] return store, contents def _read_under_pipeline(pipeline_path: str, store: MemoryStore) -> Any: """Re-open an existing store under the chosen pipeline and read it whole.""" - with zarr_config.set({"codec_pipeline.path": pipeline_path}): - with _ignore_sharding_combo_warning(): - arr = zarr.open_array(store=store, mode="r") - return arr[...] + with _ignore_sharding_combo_warning(): + arr = zarr.open_array(store=store, mode="r") + return _under_pipeline(arr, pipeline_path)[...] _BATCHED = "zarr.core.codec_pipeline.BatchedCodecPipeline" @@ -407,20 +426,20 @@ def test_pipeline_parity_subchunk_write_order( def run(pipeline_path: str) -> tuple[dict[str, bytes], Any]: store = MemoryStore() - with zarr_config.set({"codec_pipeline.path": pipeline_path}): - arr = zarr.create_array( - store=store, - shape=shape, - chunks=shard_shape, # array "chunks" == shard size for a ShardingCodec serializer - dtype="int32", - fill_value=-1, - serializer=serializer, - compressors=None, - config={"write_empty_chunks": True}, - ) - arr[:] = ref # dense full write - arr[3:9, 1:6] = 777 # partial write INTO the dense shard - contents = arr[...] + arr = zarr.create_array( + store=store, + shape=shape, + chunks=shard_shape, # array "chunks" == shard size for a ShardingCodec serializer + dtype="int32", + fill_value=-1, + serializer=serializer, + compressors=None, + config={"write_empty_chunks": True}, + ) + arr = _under_pipeline(arr, pipeline_path) + arr[:] = ref # dense full write + arr[3:9, 1:6] = 777 # partial write INTO the dense shard + contents = arr[...] return _store_snapshot(store), contents batched_bytes, batched_contents = run(_BATCHED) @@ -502,27 +521,27 @@ def test_pipeline_parity_outer_aa_bb_codecs( data = (np.arange(int(np.prod(shape))).reshape(shape) + 1).astype("uint16") store = MemoryStore() - with zarr_config.set({"codec_pipeline.path": writer}): - with _ignore_sharding_combo_warning(): - arr = zarr.create_array( - store=store, - shape=shape, - chunks=(4, 4), - dtype=data.dtype, - fill_value=0, - **codec_kwargs, - ) - arr[...] = data - arr[2:5, 1:3] = 99 # region write -- exercises the partial-encode branch + with _ignore_sharding_combo_warning(): + arr = zarr.create_array( + store=store, + shape=shape, + chunks=(4, 4), + dtype=data.dtype, + fill_value=0, + **codec_kwargs, + ) + arr = _under_pipeline(arr, writer) + arr[...] = data + arr[2:5, 1:3] = 99 # region write -- exercises the partial-encode branch expected = data.copy() expected[2:5, 1:3] = 99 - with zarr_config.set({"codec_pipeline.path": reader}): - with _ignore_sharding_combo_warning(): - arr2 = zarr.open_array(store=store, mode="r") - full = arr2[...] - partial = arr2[1:3, 2:7] # region read -- exercises the partial-decode branch + with _ignore_sharding_combo_warning(): + arr2 = zarr.open_array(store=store, mode="r") + arr2 = _under_pipeline(arr2, reader) + full = arr2[...] + partial = arr2[1:3, 2:7] # region read -- exercises the partial-decode branch np.testing.assert_array_equal(full, expected) np.testing.assert_array_equal(partial, expected[1:3, 2:7])