diff --git a/.github/workflows/zarr-codec-extraction.yml b/.github/workflows/zarr-codec-extraction.yml new file mode 100644 index 0000000000..2bc491892d --- /dev/null +++ b/.github/workflows/zarr-codec-extraction.yml @@ -0,0 +1,50 @@ +name: zarr-codec extraction + +on: + pull_request: + paths: + - 'packages/zarr-codec/**' + - 'src/zarr/abc/codec.py' + - '.github/workflows/zarr-codec-extraction.yml' + push: + branches: [main] + paths: + - 'packages/zarr-codec/**' + - 'src/zarr/abc/codec.py' + - '.github/workflows/zarr-codec-extraction.yml' + workflow_dispatch: + +permissions: + contents: read + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ['3.12', '3.14'] + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + - name: Install Python and Hatch + run: | + uv python install ${{ matrix.python-version }} + uv tool install hatch==1.18.0 + - name: Test extracted interfaces against this checkout + env: + PYTHONPATH: packages/zarr-codec/src + run: >- + hatch run test.py${{ matrix.python-version }}-minimal:pytest + packages/zarr-codec/tests --import-mode=importlib + - name: Build codec distributions + working-directory: packages/zarr-codec + run: hatch build + - name: Test built wheels + env: + PYTHONPATH: packages/zarr-codec/dist/zarr_codec-0.1.0-py3-none-any.whl + run: >- + hatch run test.py${{ matrix.python-version }}-minimal:pytest + packages/zarr-codec/tests --import-mode=importlib diff --git a/packages/zarr-codec/CHANGELOG.md b/packages/zarr-codec/CHANGELOG.md new file mode 100644 index 0000000000..9e40c8658b --- /dev/null +++ b/packages/zarr-codec/CHANGELOG.md @@ -0,0 +1,5 @@ +# zarr-codec changelog + +## Unreleased + +- Extract the existing interface into `zarr_codec.legacy` without changing Zarr runtime imports. diff --git a/packages/zarr-codec/LICENSE.txt b/packages/zarr-codec/LICENSE.txt new file mode 100644 index 0000000000..1e8da4d242 --- /dev/null +++ b/packages/zarr-codec/LICENSE.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2015-2025 Zarr Developers + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/packages/zarr-codec/README.md b/packages/zarr-codec/README.md new file mode 100644 index 0000000000..e4bd5a6eb1 --- /dev/null +++ b/packages/zarr-codec/README.md @@ -0,0 +1,54 @@ +# zarr-codec + +An initial extraction of the existing Zarr-Python codec interfaces: +codec base classes, partial-IO mixins, sync capability protocols, and CodecPipeline. + +```python +from zarr_codec.legacy import BytesBytesCodec +``` + +## Status and scope + +This is an experimental extraction draft, not a replacement Zarr backend. +The `legacy` namespace reproduces `src/zarr/abc/codec.py` from Zarr-Python commit +`1b16efee6`. Signatures, docstrings, inherited implementations, and async +behavior are preserved. New APIs can be developed under a different namespace; +none are introduced here. Built-in concrete stores/codecs remain in `zarr`. + +**There is still a runtime dependency on `zarr`.** The legacy interfaces use +its shared foundations. This package must not become a dependency of `zarr` +until that dependency is removed. The tested compatibility baseline is the +source checkout at the commit above; the dependency range is not a claim that +every release in that range has been tested. + +These are independent class definitions, not aliases to Zarr's classes. +Existing Zarr entry points do not yet recognize them as their own nominal +base classes. Do not switch an existing extension to this namespace and expect +it to plug into Zarr before runtime integration lands. At that point, the old +Zarr import paths should re-export one canonical definition, preserving class +identity. There are no deprecations or changes to Zarr's runtime in this draft. + +## Remaining extraction dependencies + +- Metadata and its recursive serialization behavior. +- Buffer and NDBuffer, including runtime generic type bounds. +- NamedConfig, configuration, and concurrent_map for batching. +- Annotation dependencies: ArraySpec, dtype classes, metadata, indexing, and store interfaces. + +The global config remains shared with Zarr, preserving concurrency behavior. +Replacing these dependencies with smaller protocols would be an API design +change and is deliberately deferred. + +## Development + +From the repository root, expose this draft package and use a Hatch test environment: + +```sh +PYTHONPATH=packages/zarr-codec/src hatch run test.py3.12-minimal:pytest packages/zarr-codec/tests --import-mode=importlib +``` + +The contract tests compare the draft against the source checkout. Behavioral +tests exercise third-party-style subclasses, not only copied signatures. + +Build this distribution from its directory with `hatch build`. API documentation +is in `docs/api/index.md`; inherited docstrings remain in the extracted module. diff --git a/packages/zarr-codec/docs/api/index.md b/packages/zarr-codec/docs/api/index.md new file mode 100644 index 0000000000..47f8cdd4a7 --- /dev/null +++ b/packages/zarr-codec/docs/api/index.md @@ -0,0 +1,3 @@ +# Legacy codec API + +::: zarr_codec.legacy diff --git a/packages/zarr-codec/pyproject.toml b/packages/zarr-codec/pyproject.toml new file mode 100644 index 0000000000..5b62bdfb73 --- /dev/null +++ b/packages/zarr-codec/pyproject.toml @@ -0,0 +1,38 @@ +[build-system] +requires = ["hatchling>=1.29.0"] +build-backend = "hatchling.build" + +[project] +name = "zarr-codec" +version = "0.1.0" +description = "Legacy Zarr codec interfaces, extracted for API evolution." +readme = "README.md" +requires-python = ">=3.12" +license = "MIT" +license-files = ["LICENSE.txt"] +authors = [{ name = "Davis Bennett", email = "davis.v.bennett@gmail.com" }] +# Temporary: legacy interfaces still consume Zarr's shared foundations. +# Remove this dependency before making zarr depend on this distribution. +dependencies = ["zarr>=3.3,<3.4", "typing_extensions>=4.14"] + +[project.urls] +Source = "https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-codec" +Issues = "https://github.com/zarr-developers/zarr-python/issues" + +[tool.hatch.build.targets.wheel] +packages = ["src/zarr_codec"] + +[tool.hatch.build.targets.sdist] +include = ["/src", "/tests", "/docs", "/CHANGELOG.md"] + +[tool.pytest.ini_options] +addopts = ["--import-mode=importlib"] +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" + +[tool.ruff] +extend = "../../pyproject.toml" + +# Match the parent project's validation of the verbatim legacy docstrings. +[tool.numpydoc_validation] +checks = ["GL10", "SS04", "PR02", "PR03", "PR05", "PR06"] diff --git a/packages/zarr-codec/src/zarr_codec/__init__.py b/packages/zarr-codec/src/zarr_codec/__init__.py new file mode 100644 index 0000000000..5f56c5b91e --- /dev/null +++ b/packages/zarr-codec/src/zarr_codec/__init__.py @@ -0,0 +1,5 @@ +"""Extracted Zarr codec interfaces; see :mod:`zarr_codec.legacy`.""" + +from zarr_codec import legacy + +__all__ = ["legacy"] diff --git a/packages/zarr-codec/src/zarr_codec/legacy.py b/packages/zarr-codec/src/zarr_codec/legacy.py new file mode 100644 index 0000000000..ce7a8d9aa7 --- /dev/null +++ b/packages/zarr-codec/src/zarr_codec/legacy.py @@ -0,0 +1,529 @@ +from __future__ import annotations + +from abc import abstractmethod +from collections.abc import Mapping +from typing import TYPE_CHECKING, Literal, Protocol, TypeGuard, runtime_checkable + +from typing_extensions import ReadOnly, TypedDict +from zarr.abc.metadata import Metadata +from zarr.core.buffer import Buffer, NDBuffer +from zarr.core.common import NamedConfig, concurrent_map +from zarr.core.config import config + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable, Iterable + from typing import Self + + from zarr.abc.store import ByteGetter, ByteSetter, Store + from zarr.core.array_spec import ArraySpec + from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType + from zarr.core.indexing import SelectorTuple + from zarr.core.metadata import ArrayMetadata + from zarr.core.metadata.v3 import ChunkGridMetadata + +__all__ = [ + "ArrayArrayCodec", + "ArrayBytesCodec", + "ArrayBytesCodecPartialDecodeMixin", + "ArrayBytesCodecPartialEncodeMixin", + "BaseCodec", + "BytesBytesCodec", + "CodecInput", + "CodecOutput", + "CodecPipeline", + "GetResult", + "SupportsSyncCodec", +] + + +class GetResult(TypedDict): + """Metadata about a store get operation.""" + + status: Literal["present", "missing"] + + +type CodecInput = NDBuffer | Buffer +type CodecOutput = NDBuffer | Buffer + + +class CodecJSON_V2[TName: str](TypedDict): + """The JSON representation of a codec for Zarr V2""" + + id: ReadOnly[TName] + + +def _check_codecjson_v2(data: object) -> TypeGuard[CodecJSON_V2[str]]: + return isinstance(data, Mapping) and "id" in data and isinstance(data["id"], str) + + +CodecJSON_V3 = str | NamedConfig[str, Mapping[str, object]] +"""The JSON representation of a codec for Zarr V3.""" + +# The widest type we will *accept* for a codec JSON +# This covers v2 and v3 +CodecJSON = str | Mapping[str, object] +"""The widest type of JSON-like input that could specify a codec.""" + + +@runtime_checkable +class SupportsSyncCodec[CI: CodecInput, CO: CodecOutput](Protocol): + """Protocol for codecs that support synchronous encode/decode. + + Codecs implementing this protocol provide `_decode_sync` and `_encode_sync` + methods that perform encoding/decoding without requiring an async event loop. + + The type parameters mirror `BaseCodec`: `CI` is the decoded type and `CO` is + the encoded type. + """ + + def _decode_sync(self, chunk_data: CO, chunk_spec: ArraySpec) -> CI: ... + + def _encode_sync(self, chunk_data: CI, chunk_spec: ArraySpec) -> CO | None: ... + + +def _codec_supports_sync(codec: object) -> bool: + """Whether `codec` can actually run on a synchronous (no event loop) path. + + Structural membership in `SupportsSyncCodec` is necessary but not always + sufficient: a codec can provide `_decode_sync`/`_encode_sync` whose ability + to run depends on runtime configuration the type system cannot see. + `ShardingCodec` is the canonical case — its sync methods delegate to its + configured inner and index codec chains, so they only work when every codec + in those chains is itself sync-capable. Such codecs opt out dynamically via + a `_sync_capable` attribute/property (absent means capable). + """ + return isinstance(codec, SupportsSyncCodec) and getattr(codec, "_sync_capable", True) + + +class BaseCodec[CI: CodecInput, CO: CodecOutput](Metadata): + """Generic base class for codecs. + + Codecs can be registered via zarr.codecs.registry. + + Warnings + -------- + This class is not intended to be directly, please use + ArrayArrayCodec, ArrayBytesCodec or BytesBytesCodec for subclassing. + """ + + # Whether this codec's encoded output is a fixed size given a fixed input + # size. Defaults to False (the conservative answer): a codec that does not + # explicitly opt in is treated as variable-size, which only disables + # size-dependent fast paths (e.g. the sharding bulk-decode), never + # correctness. Codecs with genuinely fixed-size output (BytesCodec, + # TransposeCodec, ...) override this with True. The default also keeps + # third-party / variable-length codecs (VLenUTF8, numcodecs wrappers) that + # never set the attribute from raising AttributeError where it is read. + is_fixed_size: bool = False + + @abstractmethod + def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int: + """Given an input byte length, this method returns the output byte length. + Raises a NotImplementedError for codecs with variable-sized outputs (e.g. compressors). + + Parameters + ---------- + input_byte_length : int + chunk_spec : ArraySpec + + Returns + ------- + int + """ + ... + + def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: + """Computed the spec of the chunk after it has been encoded by the codec. + This is important for codecs that change the shape, data type or fill value of a chunk. + The spec will then be used for subsequent codecs in the pipeline. + + Parameters + ---------- + chunk_spec : ArraySpec + + Returns + ------- + ArraySpec + """ + return chunk_spec + + def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: + """Fills in codec configuration parameters that can be automatically + inferred from the array metadata. + + Parameters + ---------- + array_spec : ArraySpec + + Returns + ------- + Self + """ + return self + + def validate( + self, + *, + shape: tuple[int, ...], + dtype: ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGridMetadata, + ) -> None: + """Validates that the codec configuration is compatible with the array metadata. + Raises errors when the codec configuration is not compatible. + + Parameters + ---------- + shape : tuple[int, ...] + The array shape + dtype : np.dtype[Any] + The array data type + chunk_grid : ChunkGridMetadata + The array chunk grid metadata + """ + + async def _decode_single(self, chunk_data: CO, chunk_spec: ArraySpec) -> CI: + raise NotImplementedError # pragma: no cover + + async def decode( + self, + chunks_and_specs: Iterable[tuple[CO | None, ArraySpec]], + ) -> Iterable[CI | None]: + """Decodes a batch of chunks. + Chunks can be None in which case they are ignored by the codec. + + Parameters + ---------- + chunks_and_specs : Iterable[tuple[CodecOutput | None, ArraySpec]] + Ordered set of encoded chunks with their accompanying chunk spec. + + Returns + ------- + Iterable[CI | None] + """ + return await _batching_helper(self._decode_single, chunks_and_specs) + + async def _encode_single(self, chunk_data: CI, chunk_spec: ArraySpec) -> CO | None: + raise NotImplementedError # pragma: no cover + + async def encode( + self, + chunks_and_specs: Iterable[tuple[CI | None, ArraySpec]], + ) -> Iterable[CO | None]: + """Encodes a batch of chunks. + Chunks can be None in which case they are ignored by the codec. + + Parameters + ---------- + chunks_and_specs : Iterable[tuple[CI | None, ArraySpec]] + Ordered set of to-be-encoded chunks with their accompanying chunk spec. + + Returns + ------- + Iterable[CodecOutput | None] + """ + return await _batching_helper(self._encode_single, chunks_and_specs) + + +class ArrayArrayCodec(BaseCodec[NDBuffer, NDBuffer]): + """Base class for array-to-array codecs.""" + + +class ArrayBytesCodec(BaseCodec[NDBuffer, Buffer]): + """Base class for array-to-bytes codecs.""" + + +class BytesBytesCodec(BaseCodec[Buffer, Buffer]): + """Base class for bytes-to-bytes codecs.""" + + +Codec = ArrayArrayCodec | ArrayBytesCodec | BytesBytesCodec + + +class ArrayBytesCodecPartialDecodeMixin: + """Mixin for array-to-bytes codecs that implement partial decoding.""" + + async def _decode_partial_single( + self, byte_getter: ByteGetter, selection: SelectorTuple, chunk_spec: ArraySpec + ) -> NDBuffer | None: + raise NotImplementedError + + async def decode_partial( + self, + batch_info: Iterable[tuple[ByteGetter, SelectorTuple, ArraySpec]], + ) -> Iterable[NDBuffer | None]: + """Partially decodes a batch of chunks. + This method determines parts of a chunk from the slice selection, + fetches these parts from the store (via ByteGetter) and decodes them. + + Parameters + ---------- + batch_info : Iterable[tuple[ByteGetter, SelectorTuple, ArraySpec]] + Ordered set of information about slices of encoded chunks. + The slice selection determines which parts of the chunk will be fetched. + The ByteGetter is used to fetch the necessary bytes. + The chunk spec contains information about the construction of an array from the bytes. + + Returns + ------- + Iterable[NDBuffer | None] + """ + return await concurrent_map( + list(batch_info), + self._decode_partial_single, + config.get("async.concurrency"), + ) + + +class ArrayBytesCodecPartialEncodeMixin: + """Mixin for array-to-bytes codecs that implement partial encoding.""" + + async def _encode_partial_single( + self, + byte_setter: ByteSetter, + chunk_array: NDBuffer, + selection: SelectorTuple, + chunk_spec: ArraySpec, + ) -> None: + raise NotImplementedError # pragma: no cover + + async def encode_partial( + self, + batch_info: Iterable[tuple[ByteSetter, NDBuffer, SelectorTuple, ArraySpec]], + ) -> None: + """Partially encodes a batch of chunks. + This method determines parts of a chunk from the slice selection, encodes them and + writes these parts to the store (via ByteSetter). + If merging with existing chunk data in the store is necessary, this method will + read from the store first and perform the merge. + + Parameters + ---------- + batch_info : Iterable[tuple[ByteSetter, NDBuffer, SelectorTuple, ArraySpec]] + Ordered set of information about slices of to-be-encoded chunks. + The slice selection determines which parts of the chunk will be encoded. + The ByteSetter is used to write the necessary bytes and fetch bytes for existing chunk data. + The chunk spec contains information about the chunk. + """ + await concurrent_map( + list(batch_info), + self._encode_partial_single, + config.get("async.concurrency"), + ) + + +class CodecPipeline: + """Base class for implementing CodecPipeline. + A CodecPipeline implements the read and write paths for chunk data. + On the read path, it is responsible for fetching chunks from a store (via ByteGetter), + decoding them and assembling an output array. On the write path, it encodes the chunks + and writes them to a store (via ByteSetter).""" + + @abstractmethod + def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: + """Fills in codec configuration parameters that can be automatically + inferred from the array metadata. + + Parameters + ---------- + array_spec : ArraySpec + + Returns + ------- + Self + """ + ... + + @classmethod + @abstractmethod + def from_codecs(cls, codecs: Iterable[Codec]) -> Self: + """Creates a codec pipeline from an iterable of codecs. + + Parameters + ---------- + codecs : Iterable[Codec] + + Returns + ------- + Self + """ + ... + + @classmethod + def from_array_metadata_and_store(cls, array_metadata: ArrayMetadata, store: Store) -> Self: + """Creates a codec pipeline from array metadata and a store path. + + Raises NotImplementedError by default, indicating the CodecPipeline must be created with from_codecs instead. + + Parameters + ---------- + array_metadata : ArrayMetadata + store : Store + + Returns + ------- + Self + """ + raise NotImplementedError( + f"'{type(cls).__name__}' does not implement CodecPipeline.from_array_metadata_and_store." + ) + + @property + @abstractmethod + def supports_partial_decode(self) -> bool: ... + + @property + @abstractmethod + def supports_partial_encode(self) -> bool: ... + + @abstractmethod + def validate( + self, + *, + shape: tuple[int, ...], + dtype: ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGridMetadata, + ) -> None: + """Validates that all codec configurations are compatible with the array metadata. + Raises errors when a codec configuration is not compatible. + + Parameters + ---------- + shape : tuple[int, ...] + The array shape + dtype : np.dtype[Any] + The array data type + chunk_grid : ChunkGridMetadata + The array chunk grid metadata + """ + ... + + @abstractmethod + def compute_encoded_size(self, byte_length: int, array_spec: ArraySpec) -> int: + """Given an input byte length, this method returns the output byte length. + Raises a NotImplementedError for codecs with variable-sized outputs (e.g. compressors). + + Parameters + ---------- + byte_length : int + array_spec : ArraySpec + + Returns + ------- + int + """ + ... + + @abstractmethod + async def decode( + self, + chunk_bytes_and_specs: Iterable[tuple[Buffer | None, ArraySpec]], + ) -> Iterable[NDBuffer | None]: + """Decodes a batch of chunks. + Chunks can be None in which case they are ignored by the codec. + + Parameters + ---------- + chunk_bytes_and_specs : Iterable[tuple[Buffer | None, ArraySpec]] + Ordered set of encoded chunks with their accompanying chunk spec. + + Returns + ------- + Iterable[NDBuffer | None] + """ + ... + + @abstractmethod + async def encode( + self, + chunk_arrays_and_specs: Iterable[tuple[NDBuffer | None, ArraySpec]], + ) -> Iterable[Buffer | None]: + """Encodes a batch of chunks. + Chunks can be None in which case they are ignored by the codec. + + Parameters + ---------- + chunk_arrays_and_specs : Iterable[tuple[NDBuffer | None, ArraySpec]] + Ordered set of to-be-encoded chunks with their accompanying chunk spec. + + Returns + ------- + Iterable[Buffer | None] + """ + ... + + @abstractmethod + async def read( + self, + batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]], + out: NDBuffer, + drop_axes: tuple[int, ...] = (), + ) -> tuple[GetResult, ...]: + """Reads chunk data from the store, decodes it and writes it into an output array. + Partial decoding may be utilized if the codecs and stores support it. + + Parameters + ---------- + batch_info : Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]] + Ordered set of information about the chunks. + The first slice selection determines which parts of the chunk will be fetched. + The second slice selection determines where in the output array the chunk data will be written. + The ByteGetter is used to fetch the necessary bytes. + The chunk spec contains information about the construction of an array from the bytes. + + If the Store returns ``None`` for a chunk, then the chunk was not + written and the implementation must set the values of that chunk (or + ``out``) to the fill value for the array. + + out : NDBuffer + + Returns + ------- + tuple[GetResult, ...] + One result per chunk in ``batch_info``. + """ + ... + + @abstractmethod + async def write( + self, + batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]], + value: NDBuffer, + drop_axes: tuple[int, ...] = (), + ) -> None: + """Encodes chunk data and writes it to the store. + Merges with existing chunk data by reading first, if necessary. + Partial encoding may be utilized if the codecs and stores support it. + + Parameters + ---------- + batch_info : Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]] + Ordered set of information about the chunks. + The first slice selection determines which parts of the chunk will be encoded. + The second slice selection determines where in the value array the chunk data is located. + The ByteSetter is used to fetch and write the necessary bytes. + The chunk spec contains information about the chunk. + value : NDBuffer + """ + ... + + +async def _batching_helper[CI: CodecInput, CO: CodecOutput]( + func: Callable[[CI, ArraySpec], Awaitable[CO | None]], + batch_info: Iterable[tuple[CI | None, ArraySpec]], +) -> list[CO | None]: + return await concurrent_map( + list(batch_info), + _noop_for_none(func), + config.get("async.concurrency"), + ) + + +def _noop_for_none[CI: CodecInput, CO: CodecOutput]( + func: Callable[[CI, ArraySpec], Awaitable[CO | None]], +) -> Callable[[CI | None, ArraySpec], Awaitable[CO | None]]: + async def wrap(chunk: CI | None, chunk_spec: ArraySpec) -> CO | None: + if chunk is None: + return None + return await func(chunk, chunk_spec) + + return wrap diff --git a/packages/zarr-codec/src/zarr_codec/py.typed b/packages/zarr-codec/src/zarr_codec/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/zarr-codec/tests/test_codec_behavior.py b/packages/zarr-codec/tests/test_codec_behavior.py new file mode 100644 index 0000000000..70349c5be6 --- /dev/null +++ b/packages/zarr-codec/tests/test_codec_behavior.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.core.buffer import Buffer, default_buffer_prototype +from zarr.core.buffer.cpu import Buffer as CpuBuffer +from zarr.core.dtype import UInt8 + +from zarr_codec.legacy import BytesBytesCodec, CodecPipeline + + +@dataclass(frozen=True) +class ReverseCodec(BytesBytesCodec): + name: str = "reverse" + + def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int: + return input_byte_length + + async def _encode_single(self, chunk_data: Buffer, chunk_spec: ArraySpec) -> Buffer: + if len(chunk_data) != chunk_spec.shape[0]: + raise ValueError("unexpected chunk size") + return CpuBuffer.from_bytes(chunk_data.to_bytes()[::-1]) + + async def _decode_single(self, chunk_data: Buffer, chunk_spec: ArraySpec) -> Buffer: + return await self._encode_single(chunk_data, chunk_spec) + + +@pytest.fixture +def spec() -> ArraySpec: + return ArraySpec((3,), UInt8(), 0, ArrayConfig("C", False), default_buffer_prototype()) + + +@pytest.mark.parametrize("values", [[], [None], [b"abc", None, b"xyz"], [None, b"abc"]]) +async def test_batch_roundtrip(values: list[bytes | None], spec: ArraySpec) -> None: + codec = ReverseCodec() + inputs = [CpuBuffer.from_bytes(value) if value is not None else None for value in values] + encoded = list(await codec.encode((value, spec) for value in inputs)) + assert [value.to_bytes() if value is not None else None for value in encoded] == [ + value[::-1] if value is not None else None for value in values + ] + decoded = await codec.decode((value, spec) for value in encoded) + assert [value.to_bytes() if value is not None else None for value in decoded] == values + assert codec.resolve_metadata(spec) is spec + assert codec.evolve_from_array_spec(spec) is codec + assert codec.to_dict() == {"name": "reverse"} + assert ReverseCodec.from_dict({"name": "other"}).name == "other" + + +async def test_batch_propagates_codec_error(spec: ArraySpec) -> None: + with pytest.raises(ValueError, match="unexpected chunk size"): + await ReverseCodec().encode([(CpuBuffer.from_bytes(b"too long"), spec)]) + + +def test_pipeline_requires_explicit_metadata_factory() -> None: + # The inherited method must still reject this construction path; it must + # not silently create an uninitialized pipeline. + with pytest.raises(NotImplementedError, match="from_array_metadata_and_store"): + CodecPipeline.from_array_metadata_and_store(None, None) # type: ignore[arg-type] diff --git a/packages/zarr-codec/tests/test_legacy_contract.py b/packages/zarr-codec/tests/test_legacy_contract.py new file mode 100644 index 0000000000..7c8c737e15 --- /dev/null +++ b/packages/zarr-codec/tests/test_legacy_contract.py @@ -0,0 +1,45 @@ +"""Check extension-author signatures against the extraction's source checkout.""" + +from __future__ import annotations + +import importlib +import importlib.util +import inspect + +import pytest +from zarr.abc import codec as original + + +def test_codec_package_exists() -> None: + assert importlib.util.find_spec("zarr_codec") is not None + + +@pytest.mark.parametrize( + "name", + [ + "BaseCodec", + "ArrayArrayCodec", + "ArrayBytesCodec", + "BytesBytesCodec", + "ArrayBytesCodecPartialDecodeMixin", + "ArrayBytesCodecPartialEncodeMixin", + "CodecPipeline", + "SupportsSyncCodec", + ], +) +def test_legacy_codec_signatures(name: str) -> None: + extracted = importlib.import_module("zarr_codec.legacy") + before = getattr(original, name) + after = getattr(extracted, name) + assert before is not after + assert str(inspect.signature(before)) == str(inspect.signature(after)) + for member_name, member in vars(before).items(): + if isinstance(member, (classmethod, staticmethod)): + member = member.__func__ + replacement = vars(after)[member_name].__func__ + elif inspect.isfunction(member): + replacement = vars(after)[member_name] + else: + continue + assert str(inspect.signature(member)) == str(inspect.signature(replacement)) + assert inspect.iscoroutinefunction(member) == inspect.iscoroutinefunction(replacement)