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])