diff --git a/changes/420.feature.md b/changes/420.feature.md new file mode 100644 index 0000000000..e9f5e42c8a --- /dev/null +++ b/changes/420.feature.md @@ -0,0 +1 @@ +Added a `context` to `zarr.ArrayIO` and `zarr.AsyncArrayIO`: a value defining the data types, codecs and codec pipeline in scope for one array. Extension names in a stored document or in creation keyword arguments resolve through the context instead of the default registries; a kind the context leaves unset falls back to the defaults, while a kind it sets is a closed world, and a name outside it fails to resolve with an error naming the scope. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 33b7a7daef..8d6f6d00cc 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -175,6 +175,7 @@ from zarr.abc.store import Store from zarr.codecs.sharding import IndexLocation, ShardingCodec from zarr.core.buffer import Buffer + from zarr.core.context import Context from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar from zarr.storage import StoreLike from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 @@ -220,9 +221,13 @@ def _chunk_sizes_from_shape( return tuple(result) -def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata: +def parse_array_metadata( + data: Any, path: str | None = None, *, context: Context | None = None +) -> ArrayMetadata: """Array metadata from a metadata object or a metadata document, naming the array at - `path` in warnings about how an invalid document was read. + `path` in warnings about how an invalid document was read. With a `context`, a + document's extension names resolve through it instead of the default registries; + a metadata object is already resolved and passes through any context. `ArrayV2Metadata` accepts a chunk size of 0, as it always has, though only an invalid document holds one: such metadata is read as the documents it would store @@ -235,7 +240,7 @@ def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata: if isinstance(data, dict): zarr_format = data.get("zarr_format") if zarr_format == 3: - meta_out = ArrayV3Metadata.from_dict(data, path=path) + meta_out = ArrayV3Metadata.from_dict(data, path=path, context=context) if len(meta_out.storage_transformers) > 0: msg = ( f"Array metadata contains storage transformers: {meta_out.storage_transformers}." @@ -244,16 +249,19 @@ def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata: raise ValueError(msg) return meta_out elif zarr_format == 2: - return ArrayV2Metadata.from_dict(data, path=path) + return ArrayV2Metadata.from_dict(data, path=path, context=context) else: raise ValueError(f"Invalid zarr_format: {zarr_format}. Expected 2 or 3") raise TypeError # pragma: no cover -def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None) -> CodecPipeline: +def create_codec_pipeline( + metadata: ArrayMetadata, *, store: Store | None = None, context: Context | None = None +) -> CodecPipeline: + pipeline_class = get_pipeline_class() if context is None else context.pipeline_class() if store is not None: try: - return get_pipeline_class().from_array_metadata_and_store( + return pipeline_class.from_array_metadata_and_store( array_metadata=metadata, store=store ) except NotImplementedError: @@ -266,7 +274,7 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None # "disables partial reads" warning) fire here; `evolve_from_array_spec` # re-splits the same already-warned-about chain via # `codecs_from_list_unchecked`, so it does not re-emit them. - pipeline = get_pipeline_class().from_codecs(metadata.codecs) + pipeline = pipeline_class.from_codecs(metadata.codecs) # Use the regular chunk shape if available, otherwise use a # placeholder. The ChunkTransform is shape-agnostic — the actual @@ -285,7 +293,7 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None return pipeline.evolve_from_array_spec(chunk_spec) elif isinstance(metadata, ArrayV2Metadata): v2_codec = V2Codec(filters=metadata.filters, compressor=metadata.compressor) - pipeline = get_pipeline_class().from_codecs([v2_codec]) + pipeline = pipeline_class.from_codecs([v2_codec]) chunk_spec = ArraySpec( shape=metadata.chunks, dtype=metadata.dtype, @@ -408,6 +416,7 @@ def __init__( metadata: ArrayV2Metadata | ArrayV2MetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> None: ... @overload @@ -416,6 +425,7 @@ def __init__( metadata: ArrayV3Metadata | ArrayMetadataJSON_V3, store_path: StorePath, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> None: ... def __init__( @@ -423,8 +433,9 @@ def __init__( metadata: ArrayMetadata | ArrayMetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> None: - object.__setattr__(self, "_io", AsyncArrayIO(metadata, store_path, config)) + object.__setattr__(self, "_io", AsyncArrayIO(metadata, store_path, config, context)) # An AsyncArray is the resolving layer: build its codec pipeline now, so # pipeline-construction warnings fire at creation, as they always have. _ = self._io._resolved_codec_pipeline @@ -4607,13 +4618,17 @@ def build_array_metadata( attributes: dict[str, JSON] | None = None, chunk_key_encoding: ChunkKeyEncodingLike | None = None, dimension_names: DimensionNamesLike = None, + context: Context | None = None, ) -> ArrayV2Metadata | ArrayV3Metadata: """Build an array metadata document from creation-style keyword arguments. This is the pure half of `init_array`: it normalizes the same keyword arguments into an in-memory metadata object without touching any store. The parameters have the same meaning as in - [`zarr.api.asynchronous.create_array`][]. + [`zarr.api.asynchronous.create_array`][]. With a `context`, the data type + and Zarr format 3 codec names resolve through it instead of the default + registries; Zarr format 2 filters and compressors are numcodecs objects, + which no context governs. """ if zarr_format is None: @@ -4621,7 +4636,7 @@ def build_array_metadata( from zarr.codecs.sharding import ShardingCodec - zdtype = parse_dtype(dtype, zarr_format=zarr_format) + zdtype = parse_dtype(dtype, zarr_format=zarr_format, context=context) shape_parsed = parse_shapelike(shape) item_size = zdtype.item_size if isinstance(zdtype, HasItemSize) else 1 chunk_key_encoding_parsed = _parse_chunk_key_encoding( @@ -4697,6 +4712,7 @@ def build_array_metadata( filters=filters, serializer=serializer, dtype=zdtype, + context=context, ) sub_codecs = cast("tuple[Codec, ...]", (*array_array, array_bytes, *bytes_bytes)) # A stored chunk grid passed as chunks= / shards= (e.g. by the "keep" @@ -5320,6 +5336,7 @@ def _parse_chunk_encoding_v3( filters: FiltersLike, serializer: SerializerLike, dtype: ZDType[TBaseDType, TBaseScalar], + context: Context | None = None, ) -> tuple[tuple[ArrayArrayCodec, ...], ArrayBytesCodec, tuple[BytesBytesCodec, ...]]: """ Generate chunk encoding classes for v3 arrays with optional defaults. @@ -5336,7 +5353,9 @@ def _parse_chunk_encoding_v3( maybe_array_array = (filters,) else: maybe_array_array = cast("Iterable[Codec | dict[str, JSON]]", filters) - out_array_array = tuple(_parse_array_array_codec(c) for c in maybe_array_array) + out_array_array = tuple( + _parse_array_array_codec(c, context=context) for c in maybe_array_array + ) if serializer == "auto": out_array_bytes = default_serializer_v3(dtype) @@ -5344,7 +5363,7 @@ def _parse_chunk_encoding_v3( # TODO: ensure that the serializer is compatible with the ndarray produced by the # array-array codecs. For example, if a sequence of array-array codecs produces an # array with a single-byte data type, then the serializer should not specify endiannesss. - out_array_bytes = _parse_array_bytes_codec(serializer) + out_array_bytes = _parse_array_bytes_codec(serializer, context=context) if compressors is None: out_bytes_bytes: tuple[BytesBytesCodec, ...] = () @@ -5357,7 +5376,9 @@ def _parse_chunk_encoding_v3( else: maybe_bytes_bytes = cast("Iterable[Codec | dict[str, JSON]]", compressors) - out_bytes_bytes = tuple(_parse_bytes_bytes_codec(c) for c in maybe_bytes_bytes) + out_bytes_bytes = tuple( + _parse_bytes_bytes_codec(c, context=context) for c in maybe_bytes_bytes + ) # TODO: ensure that the serializer is compatible with the ndarray produced by the # array-array codecs. For example, if a sequence of array-array codecs produces an @@ -6254,10 +6275,18 @@ class AsyncArrayIO: opened from a stored document that needed a repair stores that repair before its first chunk write. Call `store_metadata` first to make that write explicit. + + A `context` (see [`zarr.core.context.Context`][]) defines the data types, + codecs and codec pipeline in scope for this array. It is consumed where + the document resolves into classes — building the data path — so the + document operations never touch it: a name outside the context fails only + when chunks are read or written, even when the default registries know it. + Without one, the default registries are used. """ _store_path: StorePath _config: ArrayConfig + _context: Context | None _source: ArrayDocumentModel | ArrayMetadata _document_cache: ArrayDocumentModel | None _resolved_cache: ArrayMetadata | None @@ -6270,9 +6299,11 @@ def __init__( document: ArrayDocumentModel | ArrayMetadata | ArrayMetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> None: self._store_path = store_path self._config = parse_array_config(config) + self._context = context self._document_cache = None self._resolved_cache = None self._chunk_grid_cache = None @@ -6284,7 +6315,7 @@ def __init__( # A resolved metadata object passes through; a raw document dict is # read eagerly, as repaired when it needs that, warning as `AsyncArray` # always has. - self._source = parse_array_metadata(document, str(store_path)) + self._source = parse_array_metadata(document, str(store_path), context=context) @classmethod async def build( @@ -6306,12 +6337,15 @@ async def build( chunk_key_encoding: ChunkKeyEncodingLike | None = None, dimension_names: DimensionNamesLike = None, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> AsyncArrayIO: """Declare an array bound to a location in a store, without storing anything. The keyword arguments have the same meaning as in [`zarr.api.asynchronous.create_array`][]. The returned instance holds - the normalized metadata in memory; `store_metadata` stores it. + the normalized metadata in memory; `store_metadata` stores it. With a + `context`, the data type and codec names resolve through it instead of + the default registries. """ metadata = build_array_metadata( shape=shape, @@ -6327,9 +6361,10 @@ async def build( attributes=attributes, chunk_key_encoding=chunk_key_encoding, dimension_names=dimension_names, + context=context, ) store_path = await make_store_path(store, path=name) - return cls(metadata, store_path, config) + return cls(metadata, store_path, config, context=context) @classmethod async def open( @@ -6339,6 +6374,7 @@ async def open( name: str | None = None, zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> AsyncArrayIO: """Read the metadata document of an existing array and bind to it. @@ -6354,7 +6390,7 @@ async def open( """ store_path = await make_store_path(store, path=name) data = await get_array_metadata(store_path, zarr_format=zarr_format) - return cls(_array_document_model(data), store_path, config) + return cls(_array_document_model(data), store_path, config, context=context) @property def document(self) -> ArrayDocumentModel: @@ -6382,7 +6418,9 @@ def metadata(self) -> ArrayMetadata: data = _document_json(source) if isinstance(source, ZarrV2ArrayMetadata): data["attributes"] = dict(source.attributes) - self._resolved_cache = parse_array_metadata(data, str(self._store_path)) + self._resolved_cache = parse_array_metadata( + data, str(self._store_path), context=self._context + ) return self._resolved_cache @property @@ -6495,7 +6533,7 @@ def _shard_grid_shape(self) -> tuple[int, ...]: def _resolved_codec_pipeline(self) -> CodecPipeline: if self._codec_pipeline_cache is None: self._codec_pipeline_cache = create_codec_pipeline( - metadata=self.metadata, store=self._store_path.store + metadata=self.metadata, store=self._store_path.store, context=self._context ) return self._codec_pipeline_cache @@ -6714,8 +6752,9 @@ def __init__( document: ArrayDocumentModel | ArrayMetadata | ArrayMetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> None: - self._async_io = AsyncArrayIO(document, store_path, config) + self._async_io = AsyncArrayIO(document, store_path, config, context=context) @classmethod def _wrap(cls, async_io: AsyncArrayIO) -> ArrayIO: @@ -6743,6 +6782,7 @@ def build( chunk_key_encoding: ChunkKeyEncodingLike | None = None, dimension_names: DimensionNamesLike = None, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> ArrayIO: """Declare an array bound to a location in a store, without storing anything. @@ -6767,6 +6807,7 @@ def build( chunk_key_encoding=chunk_key_encoding, dimension_names=dimension_names, config=config, + context=context, ) ) ) @@ -6779,13 +6820,18 @@ def open( name: str | None = None, zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, + context: Context | None = None, ) -> ArrayIO: """Read the metadata document of an existing array and bind to it. See `AsyncArrayIO.open`. """ return cls._wrap( - sync(AsyncArrayIO.open(store, name=name, zarr_format=zarr_format, config=config)) + sync( + AsyncArrayIO.open( + store, name=name, zarr_format=zarr_format, config=config, context=context + ) + ) ) @property diff --git a/src/zarr/core/context.py b/src/zarr/core/context.py index 494945d9f0..e07f0a75f6 100644 --- a/src/zarr/core/context.py +++ b/src/zarr/core/context.py @@ -7,7 +7,8 @@ threaded through the read by zarr. (The pair mirrors `referencing.Registry` and `referencing.Resolver` in python-jsonschema.) -Both are internal: nothing in the public API takes a context yet. +The experimental explicit-IO array layer (`zarr.experimental.ArrayIO` and +`AsyncArrayIO`) takes a context publicly; everywhere else, both remain internal. """ from __future__ import annotations @@ -15,7 +16,12 @@ from dataclasses import dataclass, replace from typing import TYPE_CHECKING +from zarr.errors import UnknownCodecError + if TYPE_CHECKING: + from collections.abc import Mapping + + from zarr.abc.codec import Codec, CodecPipeline from zarr.core.common import ZarrFormat from zarr.core.dtype.common import DTypeJSON from zarr.core.dtype.registry import DataTypeRegistry @@ -27,13 +33,26 @@ class Context: """ The extensions that Zarr metadata is read with, by kind. + A kind that is `None` resolves from the default registries, exactly as it + would with no context; a kind that is set is a closed world, and a name + outside it fails to resolve even when the default registries know it. + Attributes ---------- data_types : DataTypeRegistry The data types a read may resolve. + codecs : Mapping[str, type[Codec]] | None + The codec class each codec name resolves to, or `None` for the + default codec registry. Codec instances are never resolved by name, + so they pass through any context. + codec_pipeline : type[CodecPipeline] | None + The codec pipeline class arrays are read and written with, or `None` + for the pipeline the global configuration selects. """ data_types: DataTypeRegistry + codecs: Mapping[str, type[Codec]] | None = None + codec_pipeline: type[CodecPipeline] | None = None @classmethod def default(cls) -> Context: @@ -48,6 +67,31 @@ def default(cls) -> Context: return cls(data_types=data_type_registry) + def codec_class(self, name: str) -> type[Codec]: + """The codec class `name` resolves to in this context. + + Raises [`zarr.errors.UnknownCodecError`][] naming the codecs in scope + when `name` is not among them. + """ + if self.codecs is None: + from zarr.registry import get_codec_class + + return get_codec_class(name) + try: + return self.codecs[name] + except KeyError: + raise UnknownCodecError( + f"Codec {name!r} is not in this context. Codecs in scope: {sorted(self.codecs)}." + ) from None + + def pipeline_class(self) -> type[CodecPipeline]: + """The codec pipeline class this context resolves to.""" + if self.codec_pipeline is None: + from zarr.registry import get_pipeline_class + + return get_pipeline_class() + return self.codec_pipeline + type JSONLocation = tuple[str | int, ...] diff --git a/src/zarr/core/dtype/__init__.py b/src/zarr/core/dtype/__init__.py index d1dbd6e2c8..91ac7d9be4 100644 --- a/src/zarr/core/dtype/__init__.py +++ b/src/zarr/core/dtype/__init__.py @@ -35,6 +35,7 @@ if TYPE_CHECKING: from zarr.core.common import ZarrFormat + from zarr.core.context import Context from zarr.core.dtype.common import DTypeJSON from collections.abc import Mapping @@ -161,9 +162,14 @@ # TODO: find a better name for this function -def get_data_type_from_native_dtype(dtype: npt.DTypeLike) -> ZDType[TBaseDType, TBaseScalar]: +def get_data_type_from_native_dtype( + dtype: npt.DTypeLike, *, context: Context | None = None +) -> ZDType[TBaseDType, TBaseScalar]: """ Get a data type wrapper (an instance of ``ZDType``) from a native data type, e.g. a numpy dtype. + + With a `context`, the match is made against the data types it holds instead of the + default data type registry. """ if not isinstance(dtype, np.dtype): na_dtype: np.dtype[np.generic] @@ -174,17 +180,22 @@ def get_data_type_from_native_dtype(dtype: npt.DTypeLike) -> ZDType[TBaseDType, na_dtype = np.dtype(dtype) else: na_dtype = dtype - return data_type_registry.match_dtype(dtype=na_dtype) + registry = data_type_registry if context is None else context.data_types + return registry.match_dtype(dtype=na_dtype) def get_data_type_from_json( - dtype_spec: DTypeJSON, *, zarr_format: ZarrFormat + dtype_spec: DTypeJSON, *, zarr_format: ZarrFormat, context: Context | None = None ) -> ZDType[TBaseDType, TBaseScalar]: """ Given a JSON representation of a data type and a Zarr format version, attempt to create a ZDType instance from the registered ZDType classes. + + With a `context`, the match is made against the data types it holds instead of the + default data type registry. """ - return data_type_registry.match_json(dtype_spec, zarr_format=zarr_format) + registry = data_type_registry if context is None else context.data_types + return registry.match_json(dtype_spec, zarr_format=zarr_format) def parse_data_type( @@ -234,6 +245,7 @@ def parse_dtype( dtype_spec: ZDTypeLike, *, zarr_format: ZarrFormat, + context: Context | None = None, ) -> ZDType[TBaseDType, TBaseScalar]: """ Convert the input as a ZDType. @@ -272,7 +284,7 @@ def parse_dtype( # First attempt to interpret the input as JSON if isinstance(dtype_spec, Mapping | str | Sequence): try: - return get_data_type_from_json(dtype_spec, zarr_format=zarr_format) + return get_data_type_from_json(dtype_spec, zarr_format=zarr_format, context=context) except ValueError: # no data type matched this JSON-like input pass @@ -282,7 +294,7 @@ def parse_dtype( return VariableLengthUTF8() # type: ignore[return-value] # otherwise, we have either a numpy dtype string, or a zarr v3 dtype string, and in either case # we can create a native dtype from it, and do the dtype inference from that - return get_data_type_from_native_dtype(dtype_spec) # type: ignore[arg-type] + return get_data_type_from_native_dtype(dtype_spec, context=context) # type: ignore[arg-type] def __getattr__(name: str) -> object: diff --git a/src/zarr/core/metadata/v2.py b/src/zarr/core/metadata/v2.py index 136267fc58..6be1455ee2 100644 --- a/src/zarr/core/metadata/v2.py +++ b/src/zarr/core/metadata/v2.py @@ -20,6 +20,7 @@ from zarr.core.buffer import Buffer, BufferPrototype from zarr.core.chunk_grids import ChunkGrid + from zarr.core.context import Context from zarr.core.dtype.wrapper import ( TBaseDType, TBaseScalar, @@ -152,10 +153,14 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: } @classmethod - def from_dict(cls, data: dict[str, Any], *, path: str | None = None) -> ArrayV2Metadata: + def from_dict( + cls, data: dict[str, Any], *, path: str | None = None, context: Context | None = None + ) -> ArrayV2Metadata: """Read a stored `.zarray` document (with its attributes). An invalid document that `zarr.core.metadata.repair` can read is read as repaired; a reading the user - must act on warns, naming the array at `path`.""" + must act on warns, naming the array at `path`. With a `context`, the document's + data type resolves through it instead of the default registry; filters and the + compressor are numcodecs objects, which no context governs.""" repaired, readings = repair_array_document(data, 2) # a new dict, because we are modifying it _data: dict[str, Any] = dict(repaired) @@ -178,7 +183,7 @@ def from_dict(cls, data: dict[str, Any], *, path: str | None = None) -> ArrayV2M "name": _data["dtype"], "object_codec_id": object_codec_id, } - dtype = get_data_type_from_json(dtype_spec, zarr_format=2) + dtype = get_data_type_from_json(dtype_spec, zarr_format=2, context=context) _data["dtype"] = dtype fill_value_encoded = _data.get("fill_value") diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 2d5d169e6b..74ebbe8c35 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -53,6 +53,7 @@ from zarr.codecs.sharding import ShardingCodec from zarr.core.buffer import Buffer, BufferPrototype from zarr.core.chunk_grids import ChunkGrid + from zarr.core.context import Context from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar from zarr.core.metadata.repair import ArrayDocument @@ -70,7 +71,7 @@ def parse_node_type_array(data: object) -> Literal["array"]: ) -def parse_codecs(data: object) -> tuple[Codec, ...]: +def parse_codecs(data: object, *, context: Context | None = None) -> tuple[Codec, ...]: out: tuple[Codec, ...] = () if not isinstance(data, Iterable): @@ -84,7 +85,11 @@ def parse_codecs(data: object) -> tuple[Codec, ...]: else: name_parsed, _ = parse_named_configuration(c, require_configuration=False) - codec_cls = get_codec_class(name_parsed) + codec_cls = ( + get_codec_class(name_parsed) + if context is None + else context.codec_class(name_parsed) + ) try: out += (codec_cls.from_dict(c),) except KeyError as e: @@ -665,11 +670,14 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: return {ZARR_JSON: json_to_buffer(self.to_dict(), prototype=prototype, indent=indent)} @classmethod - def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self: + def from_dict( + cls, data: dict[str, JSON], *, path: str | None = None, context: Context | None = None + ) -> Self: """Read a stored `zarr.json` array document. An invalid document that `zarr.core.metadata.repair` can read is read as repaired; a reading the user must act on warns, and a document the rectilinear chunks flag refuses raises, - naming the array at `path`.""" + naming the array at `path`. With a `context`, the document's data type and + codec names resolve through it instead of the default registries.""" # The flag gates what the document declares, so it is checked before repairs. chunk_grid = data.get("chunk_grid") if isinstance(chunk_grid, Mapping) and chunk_grid.get("name") == "rectilinear": @@ -691,7 +699,7 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self: data_type_json = _data.pop("data_type") if not check_dtype_spec_v3(data_type_json): raise ValueError(f"Invalid data_type: {data_type_json!r}") - data_type = get_data_type_from_json(data_type_json, zarr_format=3) + data_type = get_data_type_from_json(data_type_json, zarr_format=3, context=context) # check that the fill value is consistent with the data type try: @@ -713,7 +721,9 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self: shape=_data_typed["shape"], chunk_grid=_data_typed["chunk_grid"], # type: ignore[arg-type] chunk_key_encoding=_data_typed["chunk_key_encoding"], # type: ignore[arg-type] - codecs=_data_typed["codecs"], + # Codec names resolve here, under the context; `__init__`'s own + # `parse_codecs` passes the resolved instances through untouched. + codecs=parse_codecs(_data_typed["codecs"], context=context), # Attribute values are arbitrary JSON, so they have no field-specific # schema to validate. `__init__` checks the outer dict via `parse_attributes`. attributes=_data_typed.get("attributes", {}), # type: ignore[arg-type] diff --git a/src/zarr/experimental/__init__.py b/src/zarr/experimental/__init__.py index f7caaf96a1..097d142c54 100644 --- a/src/zarr/experimental/__init__.py +++ b/src/zarr/experimental/__init__.py @@ -1,5 +1,6 @@ """The experimental module is a site for exporting new or experimental Zarr features.""" from zarr.core.chunk_grids import ChunkGrid, ChunkSpec +from zarr.core.context import Context -__all__ = ["ChunkGrid", "ChunkSpec"] +__all__ = ["ChunkGrid", "ChunkSpec", "Context"] diff --git a/src/zarr/registry.py b/src/zarr/registry.py index 307c323153..332f645cfc 100644 --- a/src/zarr/registry.py +++ b/src/zarr/registry.py @@ -25,6 +25,7 @@ from zarr.core.buffer import Buffer, NDBuffer from zarr.core.chunk_key_encodings import ChunkKeyEncoding from zarr.core.common import JSON, ZarrFormat + from zarr.core.context import Context __all__ = [ "Registry", @@ -387,15 +388,22 @@ def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]: return selected_codec_cls -def _resolve_codec(data: dict[str, JSON]) -> Codec: +def _resolve_codec(data: dict[str, JSON], *, context: Context | None = None) -> Codec: """ Get a codec instance from a dict representation of that codec. + + With a `context`, the codec name resolves to a class through it instead of the + default codec registry. """ + name = data["name"] + codec_cls = get_codec_class(name) if context is None else context.codec_class(name) # type: ignore[arg-type] # TODO: narrow the type of the input to only those dicts that map on to codec class instances. - return get_codec_class(data["name"]).from_dict(data) # type: ignore[arg-type] + return codec_cls.from_dict(data) -def _parse_bytes_bytes_codec(data: dict[str, JSON] | Codec) -> BytesBytesCodec: +def _parse_bytes_bytes_codec( + data: dict[str, JSON] | Codec, *, context: Context | None = None +) -> BytesBytesCodec: """ Normalize the input to a `BytesBytesCodec` instance. If the input is already a `BytesBytesCodec`, it is returned as is. If the input is a dict, it @@ -404,7 +412,7 @@ def _parse_bytes_bytes_codec(data: dict[str, JSON] | Codec) -> BytesBytesCodec: from zarr.abc.codec import BytesBytesCodec if isinstance(data, dict): - result = _resolve_codec(data) + result = _resolve_codec(data, context=context) if not isinstance(result, BytesBytesCodec): msg = f"Expected a dict representation of a BytesBytesCodec; got a dict representation of a {type(result)} instead." raise TypeError(msg) @@ -415,7 +423,9 @@ def _parse_bytes_bytes_codec(data: dict[str, JSON] | Codec) -> BytesBytesCodec: return result -def _parse_array_bytes_codec(data: dict[str, JSON] | Codec) -> ArrayBytesCodec: +def _parse_array_bytes_codec( + data: dict[str, JSON] | Codec, *, context: Context | None = None +) -> ArrayBytesCodec: """ Normalize the input to a `ArrayBytesCodec` instance. If the input is already a `ArrayBytesCodec`, it is returned as is. If the input is a dict, it @@ -424,7 +434,7 @@ def _parse_array_bytes_codec(data: dict[str, JSON] | Codec) -> ArrayBytesCodec: from zarr.abc.codec import ArrayBytesCodec if isinstance(data, dict): - result = _resolve_codec(data) + result = _resolve_codec(data, context=context) if not isinstance(result, ArrayBytesCodec): msg = f"Expected a dict representation of an ArrayBytesCodec; got a dict representation of a {type(result)} instead." raise TypeError(msg) @@ -435,7 +445,9 @@ def _parse_array_bytes_codec(data: dict[str, JSON] | Codec) -> ArrayBytesCodec: return result -def _parse_array_array_codec(data: dict[str, JSON] | Codec) -> ArrayArrayCodec: +def _parse_array_array_codec( + data: dict[str, JSON] | Codec, *, context: Context | None = None +) -> ArrayArrayCodec: """ Normalize the input to a `ArrayArrayCodec` instance. If the input is already a `ArrayArrayCodec`, it is returned as is. If the input is a dict, it @@ -444,7 +456,7 @@ def _parse_array_array_codec(data: dict[str, JSON] | Codec) -> ArrayArrayCodec: from zarr.abc.codec import ArrayArrayCodec if isinstance(data, dict): - result = _resolve_codec(data) + result = _resolve_codec(data, context=context) if not isinstance(result, ArrayArrayCodec): msg = f"Expected a dict representation of an ArrayArrayCodec; got a dict representation of a {type(result)} instead." raise TypeError(msg) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 68f60c2e5f..b32bf6e104 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -10,12 +10,24 @@ from zarr import ArrayIO, AsyncArrayIO from zarr.api.asynchronous import open_array +from zarr.codecs import BytesCodec, ZstdCodec from zarr.core.buffer import default_buffer_prototype +from zarr.core.codec_pipeline import BatchedCodecPipeline +from zarr.core.dtype import UInt16, data_type_registry +from zarr.core.dtype.registry import DataTypeRegistry from zarr.errors import ArrayNotFoundError, UnknownCodecError +from zarr.experimental import Context from zarr.storage import MemoryStore if TYPE_CHECKING: - from zarr.core.common import ZarrFormat + from zarr.core.common import JSON, ZarrFormat + +ZSTD_JSON: dict[str, JSON] = {"name": "zstd", "configuration": {"level": 0, "checksum": False}} + + +class _MarkedPipeline(BatchedCodecPipeline): + """A pipeline class distinguishable from the default, for context tests.""" + METADATA_DOCUMENTS = ("zarr.json", ".zarray", ".zattrs", ".zgroup") @@ -202,3 +214,64 @@ async def test_erase_chunk_out_of_grid_raises() -> None: await aio.erase_chunk((2, 0)) with pytest.raises(IndexError, match="outside the chunk grid"): await aio.erase_chunk((0,)) + + +async def test_context_scopes_resolution() -> None: + """A context defines the data types, codecs and pipeline in scope: names + resolve through it when the data path is built, and the pipeline class is + honored.""" + store = MemoryStore() + data = np.arange(16, dtype="uint16").reshape(4, 4) + registry = DataTypeRegistry() + registry.register(UInt16._zarr_v3_name, UInt16) + context = Context( + data_types=registry, + codecs={"bytes": BytesCodec, "zstd": ZstdCodec}, + codec_pipeline=_MarkedPipeline, + ) + aio = await AsyncArrayIO.build( + store, + name="foo", + shape=data.shape, + dtype="uint16", + chunks=(2, 2), + compressors=ZSTD_JSON, + context=context, + ) + assert isinstance(aio.to_array().codec_pipeline, _MarkedPipeline) + await aio.setitem(Ellipsis, data) + await aio.store_metadata() + + opened = await AsyncArrayIO.open(store, name="foo", context=context) + np.testing.assert_array_equal(await opened.getitem(Ellipsis), data) + + +async def test_context_unknown_codec_resolution_raises() -> None: + """A stored codec name outside the context opens for document work and + fails to resolve when the data path is built, naming the scope.""" + store = MemoryStore() + aio = await AsyncArrayIO.build(store, name="foo", shape=(4,), dtype="uint8") + await aio.store_metadata() + context = Context(data_types=data_type_registry, codecs={"bytes": BytesCodec}) + opened = await AsyncArrayIO.open(store, name="foo", context=context) + await opened.verify() # document work needs no resolution + with pytest.raises(UnknownCodecError, match="Codecs in scope: \\['bytes'\\]"): + await opened.getitem(Ellipsis) + + +async def test_context_unknown_codec_on_build_raises() -> None: + context = Context(data_types=data_type_registry, codecs={"bytes": BytesCodec}) + with pytest.raises(UnknownCodecError, match="'zstd' is not in this context"): + await AsyncArrayIO.build( + MemoryStore(), + shape=(4,), + dtype="uint8", + compressors=ZSTD_JSON, + context=context, + ) + + +async def test_context_unknown_dtype_raises() -> None: + context = Context(data_types=DataTypeRegistry()) + with pytest.raises(ValueError, match="No Zarr data type found"): + await AsyncArrayIO.build(MemoryStore(), shape=(4,), dtype="uint16", context=context)