From addcc93da4c8ee5ae559868acb1f97ac6c72e9a0 Mon Sep 17 00:00:00 2001 From: Davis Bennett Date: Fri, 28 Aug 2026 23:42:39 +0200 Subject: [PATCH 1/5] docs+test: gzip MTIME documentation follow-up and fused-pipeline test style pass - #310 (#4296) * docs: clarify gzip MTIME handling in comparison helper and fused-pipeline test docs Follow-up to zarr-developers/zarr-python#4270: document that _gzip_streams_equal_except_mtime skips bytes 4-8 (the RFC 1952 MTIME field of the standard 10-byte gzip header), and update the AsyncChunkTransform section comment and test docstring to note that byte-identity is expected except for gzip's MTIME header field. Assisted-by: ClaudeCode:claude-fable-5 * test: style pass on fused-pipeline tests Hoist the imports that nearly every test re-imported locally, add a shared _make_spec helper to replace nine copies of the ArraySpec boilerplate, and convert test_async_chunk_transform_matches_sync to Expect cases where each case declares its expected byte-comparison function (exact equality, or gzip-MTIME-tolerant), removing the isinstance branch from the test body. Also add the missing read-back assertion in test_sync_write_async_read_roundtrip, which previously read into a buffer and never compared it to the written data. Assisted-by: ClaudeCode:claude-fable-5 * test: require a ZDType in _make_spec instead of a dtype string Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/codecs/gzip.py | 7 + tests/test_fused_pipeline.py | 292 ++++++++++++----------------------- 2 files changed, 102 insertions(+), 197 deletions(-) diff --git a/src/zarr/codecs/gzip.py b/src/zarr/codecs/gzip.py index 7f21872034..e5b54639f9 100644 --- a/src/zarr/codecs/gzip.py +++ b/src/zarr/codecs/gzip.py @@ -29,6 +29,13 @@ def parse_gzip_level(data: JSON) -> int: def _gzip_streams_equal_except_mtime(a: bytes, b: bytes) -> bool: + """Compare two gzip streams, ignoring the MTIME field of the header. + + Per RFC 1952 the gzip header is [magic(2)][CM(1)][FLG(1)][MTIME(4)][XFL(1)][OS(1)], + so bytes 4-8 are MTIME. The fixed offsets assume the standard 10-byte header + with no FNAME/FEXTRA/FCOMMENT flags set, which holds here because numcodecs' + ``GZip.encode`` wraps ``gzip.GzipFile`` without a filename. + """ if len(a) != len(b): return False diff --git a/tests/test_fused_pipeline.py b/tests/test_fused_pipeline.py index 87bfaf71b4..6025ae8caf 100644 --- a/tests/test_fused_pipeline.py +++ b/tests/test_fused_pipeline.py @@ -5,11 +5,13 @@ import asyncio from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any +from unittest.mock import patch import numpy as np import pytest import zarr +from tests.conftest import Expect from zarr.abc.codec import ( ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, @@ -21,8 +23,13 @@ from zarr.codecs.gzip import GzipCodec, _gzip_streams_equal_except_mtime from zarr.codecs.transpose import TransposeCodec from zarr.codecs.zstd import ZstdCodec -from zarr.core.codec_pipeline import FusedCodecPipeline +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.core.buffer import BufferPrototype, default_buffer_prototype +from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer +from zarr.core.chunk_utils import ChunkTransform, evolve_codecs +from zarr.core.codec_pipeline import AsyncChunkTransform, FusedCodecPipeline from zarr.core.config import config as zarr_config +from zarr.core.dtype import Float32, Float64, Int32, UInt8, ZDType from zarr.registry import register_codec from zarr.storage import MemoryStore, StorePath, WrapperStore from zarr.storage._utils import _normalize_byte_range_index @@ -32,8 +39,28 @@ from collections.abc import AsyncIterator, Callable, Iterable from zarr.abc.store import ByteRequest - from zarr.core.array_spec import ArraySpec - from zarr.core.buffer import Buffer, BufferPrototype, NDBuffer + from zarr.core.buffer import Buffer, NDBuffer + + +_FLOAT64 = Float64() + + +def _make_spec( + shape: tuple[int, ...], + zdtype: ZDType[Any, Any] = _FLOAT64, + fill_value: float = 0, + *, + write_empty_chunks: bool = True, + prototype: BufferPrototype | None = None, +) -> ArraySpec: + """An ArraySpec with the C-order defaults shared by the tests in this file.""" + return ArraySpec( + shape=shape, + dtype=zdtype, + fill_value=zdtype.cast_scalar(fill_value), + config=ArrayConfig(order="C", write_empty_chunks=write_empty_chunks), + prototype=prototype if prototype is not None else default_buffer_prototype(), + ) @pytest.mark.parametrize( @@ -62,10 +89,6 @@ def test_sync_api_compute_off_event_loop(monkeypatch: pytest.MonkeyPatch) -> Non and the penalty grows with codec cost (observed as "fused pipeline is slower for zstd data" under dask-style multi-threaded single-chunk reads). """ - import asyncio - - from zarr.core.chunk_utils import ChunkTransform - compute_on_loop = {"decode": False, "encode": False} calls = {"decode": 0, "encode": 0} real_decode = ChunkTransform.decode_chunk @@ -116,22 +139,10 @@ def traced_encode(self: ChunkTransform, chunk_array: Any, chunk_spec: Any) -> An def test_evolve_from_array_spec() -> None: """evolve_from_array_spec creates a sync transform.""" - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.dtype import get_data_type_from_native_dtype - pipeline = FusedCodecPipeline.from_codecs((BytesCodec(),)) assert pipeline.sync_transform is None - zdtype = get_data_type_from_native_dtype(np.dtype("float64")) - spec = ArraySpec( - shape=(100,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) - evolved = pipeline.evolve_from_array_spec(spec) + evolved = pipeline.evolve_from_array_spec(_make_spec((100,))) assert evolved.sync_transform is not None @@ -147,31 +158,20 @@ def test_evolve_from_array_spec() -> None: @pytest.mark.parametrize( - ("dtype", "shape"), + ("zdtype", "shape"), [ - ("float64", (100,)), - ("float32", (50,)), - ("int32", (200,)), - ("float64", (10, 10)), + (Float64(), (100,)), + (Float32(), (50,)), + (Int32(), (200,)), + (Float64(), (10, 10)), ], ids=["f64-1d", "f32-1d", "i32-1d", "f64-2d"], ) -def test_read_write_sync_roundtrip(dtype: str, shape: tuple[int, ...]) -> None: +def test_read_write_sync_roundtrip(zdtype: ZDType[Any, Any], shape: tuple[int, ...]) -> None: """Data written via write_sync can be read back via read_sync.""" - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.dtype import get_data_type_from_native_dtype - store = MemoryStore() - zdtype = get_data_type_from_native_dtype(np.dtype(dtype)) - spec = ArraySpec( - shape=shape, - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + dtype = zdtype.to_native_dtype() + spec = _make_spec(shape, zdtype) pipeline = FusedCodecPipeline.from_codecs((BytesCodec(),)) pipeline = pipeline.evolve_from_array_spec(spec) @@ -200,20 +200,8 @@ def test_read_write_sync_roundtrip(dtype: str, shape: tuple[int, ...]) -> None: def test_read_sync_missing_chunk_fills() -> None: """Sync read of a missing chunk fills with the fill value.""" - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.dtype import get_data_type_from_native_dtype - store = MemoryStore() - zdtype = get_data_type_from_native_dtype(np.dtype("float64")) - spec = ArraySpec( - shape=(10,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(42.0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + spec = _make_spec((10,), fill_value=42.0) pipeline = FusedCodecPipeline.from_codecs((BytesCodec(),)) pipeline = pipeline.evolve_from_array_spec(spec) @@ -232,21 +220,10 @@ def test_read_sync_missing_chunk_fills() -> None: def test_sync_write_async_read_roundtrip() -> None: """Data written via write_sync can be read back via async read.""" - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.dtype import get_data_type_from_native_dtype from zarr.core.sync import sync store = MemoryStore() - zdtype = get_data_type_from_native_dtype(np.dtype("float64")) - spec = ArraySpec( - shape=(100,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + spec = _make_spec((100,)) pipeline = FusedCodecPipeline.from_codecs((BytesCodec(),)) pipeline = pipeline.evolve_from_array_spec(spec) @@ -271,17 +248,14 @@ def test_sync_write_async_read_roundtrip() -> None: ) ) + np.testing.assert_array_equal(data, out.as_numpy_array()) + def test_chunk_transform_uses_runtime_prototype() -> None: """ChunkTransform must pass each codec the prototype from the runtime chunk_spec, not one captured at evolve time. Constructs ChunkTransform directly (a Fused-internal data structure with no BatchedCodecPipeline equivalent). """ - from zarr.abc.codec import BytesBytesCodec - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import BufferPrototype, default_buffer_prototype - from zarr.core.chunk_utils import ChunkTransform - from zarr.core.dtype import get_data_type_from_native_dtype class _PrototypeRecordingCodec(BytesBytesCodec): # type: ignore[misc,unused-ignore] """A no-op BB codec that records the prototype it was called with.""" @@ -319,16 +293,8 @@ async def _decode_single(self, chunk_bytes: Buffer, chunk_spec: ArraySpec) -> Bu recording = _PrototypeRecordingCodec() transform = ChunkTransform(codecs=(BytesCodec(), recording)) - zdtype = get_data_type_from_native_dtype(np.dtype("float64")) - def _spec(prototype: BufferPrototype) -> ArraySpec: - return ArraySpec( - shape=(10,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0.0), - config=ArrayConfig(order="C", write_empty_chunks=False), - prototype=prototype, - ) + return _make_spec((10,), write_empty_chunks=False, prototype=prototype) proto_default = default_buffer_prototype() # A distinct BufferPrototype instance with the same buffer/nd_buffer types -- @@ -371,8 +337,6 @@ def test_read_write_with_thread_pool() -> None: returning 1) would silently degrade all the pool tests into re-testing the sequential branch while staying green. """ - from unittest.mock import patch - import zarr.core.codec_pipeline as cp_mod with zarr_config.set(_FUSED_POOL_CONFIG): @@ -404,8 +368,6 @@ def test_read_write_with_thread_pool() -> None: def test_thread_pool_write_worker_exception_propagates() -> None: """A store error raised inside a pool worker during write_sync surfaces to the caller (write_sync consumes pool.map, so worker exceptions re-raise).""" - from unittest.mock import patch - with zarr_config.set(_FUSED_POOL_CONFIG): store = MemoryStore() arr = zarr.create_array( @@ -426,8 +388,6 @@ def test_thread_pool_write_worker_exception_propagates() -> None: def test_thread_pool_read_worker_exception_propagates() -> None: """A store error raised inside a pool worker during read_sync surfaces to the caller (read_sync consumes pool.map into a tuple).""" - from unittest.mock import patch - with zarr_config.set(_FUSED_POOL_CONFIG): store = MemoryStore() arr = zarr.create_array( @@ -474,12 +434,7 @@ async def test_encode_and_write_as_completed_cancels_stray_writes_on_failure() - never retrieved (an unraisable "Task exception was never retrieved" warning if it later fails). """ - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.chunk_utils import ChunkTransform from zarr.core.codec_pipeline import _encode_and_write_as_completed - from zarr.core.dtype import get_data_type_from_native_dtype write_started = asyncio.Event() write_finished = False @@ -517,14 +472,7 @@ async def delete(self) -> None: async def set_if_not_exists(self, default: Buffer) -> None: pass - zdtype = get_data_type_from_native_dtype(np.dtype("uint8")) - chunk_spec = ArraySpec( - shape=(1,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + chunk_spec = _make_spec((1,), UInt8()) chunk_array = CPUNDBuffer.from_numpy_array(np.zeros(1, dtype="uint8")) transform = ChunkTransform(codecs=(BytesCodec(),)) @@ -598,28 +546,12 @@ def test_shared_transform_decode_alternating_specs() -> None: underpins that guarantee. (The concurrent counterpart is `test_concurrent_reads_shared_transform_with_pool`.) """ - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.chunk_utils import ChunkTransform - from zarr.core.dtype import get_data_type_from_native_dtype - - def _spec(shape: tuple[int, ...]) -> ArraySpec: - zdtype = get_data_type_from_native_dtype(np.dtype("int32")) - return ArraySpec( - shape=shape, - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) - transform = ChunkTransform(codecs=(TransposeCodec(order=(1, 0)), BytesCodec())) # two distinct specs (different shapes) sharing the one transform + cache slot cases = [] for shape in [(5, 7), (3, 11)]: - spec = _spec(shape) + spec = _make_spec(shape, Int32()) arr = np.arange(int(np.prod(shape)), dtype="int32").reshape(shape) encoded = transform.encode_chunk(CPUNDBuffer.from_numpy_array(arr), spec) assert encoded is not None @@ -643,11 +575,6 @@ def test_sharded_fallback_inner_chunks_avoid_async_transform() -> None: dict lookup plus an async per-chunk transform — measured at 1.5x (raw) to 3.6x (gzip) of sharded fallback read time. """ - from unittest.mock import patch - - from zarr.core.codec_pipeline import AsyncChunkTransform - from zarr.testing.store import LatencyStore - calls = {"decode": 0, "encode": 0} orig_decode = AsyncChunkTransform.decode_chunk orig_encode = AsyncChunkTransform.encode_chunk @@ -716,22 +643,9 @@ def test_write_over_sync_byte_setter_takes_sync_path() -> None: directly: without it, this write degrades to the async fallback (one coroutine per inner chunk for an in-memory dict store). """ - import asyncio - from unittest.mock import patch - from zarr.codecs.sharding import _ShardingByteSetter - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.dtype import get_data_type_from_native_dtype - zdtype = get_data_type_from_native_dtype(np.dtype("uint8")) - spec = ArraySpec( - shape=(10,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + spec = _make_spec((10,), UInt8()) pipeline = FusedCodecPipeline.from_codecs([BytesCodec()]).evolve_from_array_spec(spec) assert pipeline.sync_transform is not None @@ -839,8 +753,6 @@ def test_sharded_roundtrip_with_async_only_inner_codec() -> None: # The stored bytes are valid for the default pipeline too: read them back # under BatchedCodecPipeline (default codec_pipeline.path). Opening from # metadata needs the codec name in the registry. - from zarr.registry import register_codec - register_codec("test-async-only-noop", _AsyncOnlyNoopCodec) reread = zarr.open_array(store=store, mode="r") np.testing.assert_array_equal(reread[:], data) @@ -849,51 +761,65 @@ def test_sharded_roundtrip_with_async_only_inner_codec() -> None: # --------------------------------------------------------------------------- # AsyncChunkTransform: the async per-chunk codec chain used on the async # fallback path. It is the async mirror of ChunkTransform, so it must produce -# identical bytes/arrays. The default (Fused, sync-store) path never uses it; +# identical bytes/arrays — except for gzip's embedded 4-byte MTIME header +# field, which the gzip case's comparator deliberately ignores. +# The default (Fused, sync-store) path never uses it; # these tests drive it directly over multi-codec chains so the aa/bb loops and # the all-fill drop branch are exercised. # --------------------------------------------------------------------------- -@pytest.mark.parametrize( - "codecs", - [ - (BytesCodec(),), - (BytesCodec(), GzipCodec(level=1)), - (TransposeCodec(order=(1, 0)), BytesCodec()), - (TransposeCodec(order=(1, 0)), BytesCodec(), ZstdCodec(level=1)), - ], - ids=["bytes-only", "bb", "aa", "aa+ab+bb"], -) -def test_async_chunk_transform_matches_sync(codecs: tuple[Any, ...]) -> None: +def _bytes_identical(a: bytes, b: bytes) -> bool: + return a == b + + +# Each case pairs a codec chain with the expected relationship between the +# async and sync encoders' output: byte-identical, except for a chain +# containing gzip, whose embedded 4-byte MTIME header field varies with the +# encode time. +_ASYNC_SYNC_PARITY_CASES: list[Expect[tuple[Any, ...], Callable[[bytes, bytes], bool]]] = [ + Expect( + input=(BytesCodec(),), + output=_bytes_identical, + id="bytes-only", + ), + Expect( + input=(BytesCodec(), GzipCodec(level=1)), + output=_gzip_streams_equal_except_mtime, + id="bb", + ), + Expect( + input=(TransposeCodec(order=(1, 0)), BytesCodec()), + output=_bytes_identical, + id="aa", + ), + Expect( + input=(TransposeCodec(order=(1, 0)), BytesCodec(), ZstdCodec(level=1)), + output=_bytes_identical, + id="aa+ab+bb", + ), +] + + +@pytest.mark.parametrize("case", _ASYNC_SYNC_PARITY_CASES, ids=lambda c: c.id) +def test_async_chunk_transform_matches_sync( + case: Expect[tuple[Any, ...], Callable[[bytes, bytes], bool]], +) -> None: """`AsyncChunkTransform.decode_chunk`/`encode_chunk` must round-trip and - produce exactly what the synchronous `ChunkTransform` produces, across - array->array, array->bytes, and bytes->bytes codec combinations. + produce what the synchronous `ChunkTransform` produces, across + array->array, array->bytes, and bytes->bytes codec combinations. Each + case's expected output is a byte-comparison function: exact equality, + except for gzip, whose embedded 4-byte MTIME header field is deliberately + ignored. This is the async mirror of the codecs the default pipeline runs synchronously; a divergence here corrupts data only on the async fallback path (remote stores), which no end-to-end test of the default pipeline touches. """ - import asyncio - - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.chunk_utils import ChunkTransform, evolve_codecs - from zarr.core.codec_pipeline import AsyncChunkTransform - from zarr.core.dtype import get_data_type_from_native_dtype - shape = (4, 4) - zdtype = get_data_type_from_native_dtype(np.dtype("int32")) - spec = ArraySpec( - shape=shape, - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) - evolved = evolve_codecs(codecs, spec) + spec = _make_spec(shape, Int32()) + evolved = evolve_codecs(case.input, spec) sync_t = ChunkTransform(codecs=evolved) async_t = AsyncChunkTransform(codecs=evolved) @@ -904,19 +830,7 @@ def test_async_chunk_transform_matches_sync(codecs: tuple[Any, ...]) -> None: async_bytes = asyncio.run(async_t.encode_chunk(value, spec)) assert sync_bytes is not None assert async_bytes is not None - - has_timestamp_codec = any(isinstance(c, GzipCodec) for c in evolved) - - if has_timestamp_codec: - assert _gzip_streams_equal_except_mtime( - async_bytes.to_bytes(), - sync_bytes.to_bytes(), - ) - else: - np.testing.assert_array_equal( - async_bytes.to_bytes(), - sync_bytes.to_bytes(), - ) + assert case.output(async_bytes.to_bytes(), sync_bytes.to_bytes()) sync_arr = sync_t.decode_chunk(async_bytes, spec) async_arr = asyncio.run(async_t.decode_chunk(async_bytes, spec)) @@ -928,21 +842,7 @@ def test_async_decode_encode_passes_through_none_chunks() -> None: """`FusedCodecPipeline.decode`/`encode` (the async batch entry points used on the fallback path) map a None chunk to None and leave real chunks untouched — pins the None-passthrough branch the default sync path skips.""" - import asyncio - - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.dtype import get_data_type_from_native_dtype - - zdtype = get_data_type_from_native_dtype(np.dtype("int32")) - spec = ArraySpec( - shape=(4,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + spec = _make_spec((4,), Int32()) pipeline = FusedCodecPipeline.from_codecs([BytesCodec()]).evolve_from_array_spec(spec) data = np.arange(4, dtype="int32") @@ -1199,8 +1099,6 @@ def test_sync_io_capability_gates_fused_paths( (read-modify-write) writes, and all-fill (delete) writes — a store with a partial sync surface must get a clean async fallback, never a mid-batch error.""" - from unittest.mock import patch - store = store_factory() assert _store_supports_sync_io(store) is expect_sync_path From 5b5f3a3e690c9e74cebf86952606c8bb21ac2dfa Mon Sep 17 00:00:00 2001 From: Srijan Keshri Date: Sun, 30 Aug 2026 20:01:30 +0530 Subject: [PATCH 2/5] feat: name packages that provide a codec zarr cannot find (#4277) * feat: name packages that provide a codec zarr cannot find Closes #4271. When zarr fails to resolve a codec it now says which Python packages are known to provide it, instead of raising a bare KeyError holding only the codec name: An implementation for codec 'wavpack' is not available. Register one explicitly using the codec registry (see ), or install a Python package that registers a codec implementation with numcodecs. Known packages supporting this codec: wavpack-numcodecs. Two hand-maintained tables in zarr/registry.py hold the mapping, one per Zarr format, because the two formats resolve codecs through different registries and the same name can mean different things in each: `imagecodecs_*` names are declared by `virtual-tiff` under the `zarr.codecs` entry point group and by `imagecodecs-numcodecs` under `numcodecs.codecs`, and `crc32c` is a codec zarr implements itself in format 3 while in format 2 it needs `numcodecs[crc32c]`. Each table has an exact-match and a prefix-match half, since packages that provide many codecs namespace them behind a shared prefix. Entries cover third-party packages and the codecs numcodecs gates behind its own optional dependencies -- `zfpy`, `pcodec`, `crc32c` and `msgpack2` -- which are the most common missing-codec case in practice. Backwards compatibility: `get_codec_class` now raises `zarr.errors.UnknownCodecError` instead of `KeyError`, both for a codec with no registered implementation and for a codec whose configured implementation is not registered. `get_numcodec` raises it instead of the ValueError numcodecs raises for an unregistered format 2 codec id. All are subclasses of `ValueError`. Carrying the message on a `KeyError` was not an option: `KeyError.__str__` reprs its argument, so a multi-sentence message comes back quoted and escaped. `UnknownCodecError` is now exported from `zarr.errors`, since users are being told to catch it. `get_numcodec` supports numcodecs down to the declared 0.14 floor: `numcodecs.errors` only exists from 0.15.1, so the unregistered-codec check prefers that exception type where it is importable and falls back to matching the message otherwise. Signed-off-by: arcusbuilds * fix: address review feedback 1. parse_codecs converts KeyError from from_dict again. The removed try/except wrapped the whole expression, not just the registry lookup, so a codec whose from_dict indexes a malformed configuration leaked a bare KeyError out of metadata parsing. On the zarr.open fallback path that KeyError was swallowed and reported as an unrelated group error: with mode="a" it surfaced as `TypeError: open_group() got an unexpected keyword argument 'shape'`. The catch is narrow, around from_dict only, since get_codec_class now raises for the lookup half. It raises MetadataValidationError naming the codec and the missing key rather than restoring the old message, which reported the missing configuration key as though it were the codec name ("Unknown codec: 'required_option'"). 2. The config-pin branch raises BadConfigError, matching get_pipeline_class, get_buffer_class and get_ndbuffer_class, which all use it for this exact situation. This also stops migrate_v3._find_numcodecs_zarr3 misreporting a config typo as a missing numcodecs codec. 3. Three tests assumed the advertised packages were absent. Both registries are entry-point driven, so they failed in any environment with zarr-n5 or wavpack-numcodecs installed, which are the packages the messages recommend. Two fixtures now remove the specific entry for the duration of the test. Verified by installing both packages and re-running. 4. test_mapping_does_not_shadow_builtin_codecs selected on "registry is non-empty", conflating loaded-in-this-process with implemented-by-zarr. It now selects on the implementing class's module, so a lazy-loaded third-party codec cannot fail it. 5. get_numcodec's Raises section notes that numcodecs' own error propagates unchanged when data carries no string "id". 6. Dropped the `pragma: no cover` on the numcodecs < 0.15.1 fallback. The min_deps env pins numcodecs==0.14.* and runs run-coverage, so that branch is measured. Also hoisted the repeated in-function imports in tests/test_registry.py to the module level. Signed-off-by: arcusbuilds * fix: address the second review round Four defects, all found by review after the previous round was reported clean. get_numcodec no longer wraps the numcodecs call in an exception handler. Catching cannot distinguish "this id is unregistered" from "a registered codec rejected its configuration" or "a wrapper codec failed to resolve an inner codec", and it was relabelling both of the latter with the outer id plus a package hint that was wrong. Reproduced: a wrapper registered as `wavpack` whose from_config resolved a missing inner codec reported "An implementation for codec 'wavpack' is not available ... install wavpack-numcodecs", when wavpack was installed and the missing codec was something else entirely. It now performs the lookups numcodecs performs, before delegating. That also fixes reading `id` off a non-mapping input, which raised AttributeError where a ValueError used to propagate. And it removes _is_missing_numcodec_error, and with it the numcodecs <0.15.1 compatibility branch, since there is no longer an exception to classify. Note this narrows the zarr error to Mapping inputs; a duck-typed mapping now gets numcodecs' error instead, as it did before this PR. The imagecodecs_ prefix in the Zarr format 3 table pointed at virtual-tiff, which declares 15 of the 81 imagecodecs_* names under zarr.codecs; imagecodecs-numcodecs declares all 81, but under numcodecs.codecs. Users of the other 66 names were told to install a package that does not provide them. The format 3 side now lists the 15 exact names, so the rest get no hint rather than a wrong one. The format 2 side keeps the prefix, where it is correct. test_parse_codecs_converts_keyerror_from_from_dict leaked test_picky into the global codec registry, in the file that also reads that global. Tests: mutation testing showed six mutations surviving. Added the missing coverage for a registered codec rejecting its configuration, wrapper codecs resolving an inner codec by either route, non-mapping input, the imagecodecs_ over-match, and the _resolve_codec entry point. Message assertions now pin the whole string and the URL constants rather than substrings, which had allowed both documentation URLs to be replaced with wrong ones and half the message body to be deleted with every test still passing. Also corrects the docs path to src/zarr/registry.py, and rewrites the changelog to lead with the exception-type changes and to document the parse_codecs change it had omitted. Signed-off-by: arcusbuilds --------- Signed-off-by: arcusbuilds Co-authored-by: Davis Bennett --- changes/4277.feature.md | 28 +++ docs/user-guide/extending.md | 7 + src/zarr/core/metadata/v3.py | 13 +- src/zarr/errors.py | 1 + src/zarr/metadata/migrate_v3.py | 3 +- src/zarr/registry.py | 172 ++++++++++++- tests/test_codecs/test_numcodecs.py | 14 +- tests/test_config.py | 4 +- tests/test_registry.py | 374 ++++++++++++++++++++++++++++ 9 files changed, 588 insertions(+), 28 deletions(-) create mode 100644 changes/4277.feature.md create mode 100644 tests/test_registry.py diff --git a/changes/4277.feature.md b/changes/4277.feature.md new file mode 100644 index 0000000000..f5c247496d --- /dev/null +++ b/changes/4277.feature.md @@ -0,0 +1,28 @@ +`zarr.registry.get_codec_class` now raises `zarr.errors.UnknownCodecError` instead of `KeyError` +when no implementation is registered for a codec, and `zarr.core.config.BadConfigError` instead of +`KeyError` when the implementation named in `config["codecs"][name]` is not registered. +`zarr.registry.get_numcodec` raises `UnknownCodecError` instead of the `ValueError` numcodecs +raises for an unregistered Zarr format 2 codec id (`numcodecs.errors.UnknownCodecError` on +numcodecs 0.15.1 and later). All of these are subclasses of `ValueError`, so `except ValueError` +is unaffected, but `except KeyError` and `except numcodecs.errors.UnknownCodecError` are. + +These errors now name Python packages known to provide the codec, so that a user who cannot read +an array learns what to install: + +``` +An implementation for codec 'wavpack' is not available. Register one explicitly using the codec +registry (see ...), or install a Python package that registers a codec implementation with +numcodecs. Known packages supporting this codec: wavpack-numcodecs. +``` + +The tables covering this live in `src/zarr/registry.py`, one per Zarr format, and include the +codecs `numcodecs` gates behind its own optional dependencies (`zfpy`, `pcodec`, `crc32c`, +`msgpack2`). Codec authors can add their published package to them. + +A codec whose `from_dict` raises `KeyError` on a malformed configuration now surfaces as +`zarr.errors.MetadataValidationError` naming the codec and the missing key. Previously it was +reported as `UnknownCodecError: Unknown codec: ''`, presenting a configuration +key as though it were a codec name, and on the `zarr.open` path a bare `KeyError` could be +swallowed by the array-then-group fallback and reported as an unrelated group error. + +`zarr.errors.UnknownCodecError` is now exported from `zarr.errors`. diff --git a/docs/user-guide/extending.md b/docs/user-guide/extending.md index f852f9105e..507afedea7 100644 --- a/docs/user-guide/extending.md +++ b/docs/user-guide/extending.md @@ -63,6 +63,13 @@ New codecs need to have their own unique identifier. To avoid naming collisions, strongly recommended to prefix the codec identifier with a unique name. For example, the codecs from `numcodecs` are prefixed with `numcodecs.`, e.g. `numcodecs.delta`. +If someone opens an array that uses your codec without your package installed, Zarr raises +[`zarr.errors.UnknownCodecError`][] explaining how to register an implementation. Zarr also +keeps a small table of codec names and the published packages that provide them, and names +those packages in that error. Once your package is on PyPI, please open a pull request adding +it to the codec-package tables in `src/zarr/registry.py`, so that users get a message telling them +exactly what to install. + !!! note Note that the extension mechanism for the Zarr format 3 is still under development. Requirements for custom codecs including the choice of codec identifiers might diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index fc47f8fc95..988d2d369c 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -36,7 +36,7 @@ from zarr.core.dtype.common import check_dtype_spec_v3 from zarr.core.json_parse import parse_field, validate_json_value from zarr.core.metadata.common import parse_attributes -from zarr.errors import MetadataValidationError, NodeTypeValidationError, UnknownCodecError +from zarr.errors import MetadataValidationError, NodeTypeValidationError from zarr.registry import get_codec_class if TYPE_CHECKING: @@ -74,10 +74,17 @@ def parse_codecs(data: object) -> tuple[Codec, ...]: else: name_parsed, _ = parse_named_configuration(c, require_configuration=False) + codec_cls = get_codec_class(name_parsed) try: - out += (get_codec_class(name_parsed).from_dict(c),) + out += (codec_cls.from_dict(c),) except KeyError as e: - raise UnknownCodecError(f"Unknown codec: {e.args[0]!r}") from e + # A codec's `from_dict` may index its configuration directly, so a malformed + # configuration surfaces as a KeyError. Convert it: a bare KeyError escaping + # metadata parsing is swallowed by the array-then-group fallback in + # `zarr.api.asynchronous.open`, which then reports an unrelated group error. + raise MetadataValidationError( + f"Invalid configuration for codec {name_parsed!r}: missing key {e.args[0]!r}." + ) from e return out diff --git a/src/zarr/errors.py b/src/zarr/errors.py index 781bebe534..3e445de2e9 100644 --- a/src/zarr/errors.py +++ b/src/zarr/errors.py @@ -12,6 +12,7 @@ "MetadataValidationError", "NegativeStepError", "NodeTypeValidationError", + "UnknownCodecError", "UnstableSpecificationWarning", "VindexInvalidSelectionError", "ZarrDeprecationWarning", diff --git a/src/zarr/metadata/migrate_v3.py b/src/zarr/metadata/migrate_v3.py index 370af75a6d..ad177e19fb 100644 --- a/src/zarr/metadata/migrate_v3.py +++ b/src/zarr/metadata/migrate_v3.py @@ -29,6 +29,7 @@ from zarr.core.metadata.v2 import ArrayV2Metadata from zarr.core.metadata.v3 import ArrayV3Metadata, RegularChunkGridMetadata from zarr.core.sync import sync +from zarr.errors import UnknownCodecError from zarr.registry import get_codec_class from zarr.storage import StorePath from zarr.types import AnyArray @@ -273,7 +274,7 @@ def _find_numcodecs_zarr3(numcodecs_codec: numcodecs.abc.Codec) -> Codec: try: codec_v3 = get_codec_class(numcodec_name) - except KeyError as exc: + except UnknownCodecError as exc: raise ValueError( f"Couldn't find corresponding zarr.codecs.numcodecs codec for {numcodecs_codec.codec_id}" ) from exc diff --git a/src/zarr/registry.py b/src/zarr/registry.py index c2c0eb2921..a7537e5023 100644 --- a/src/zarr/registry.py +++ b/src/zarr/registry.py @@ -2,12 +2,13 @@ import warnings from collections import defaultdict +from collections.abc import Mapping from importlib.metadata import entry_points as get_entry_points from typing import TYPE_CHECKING, Any from zarr.core.config import BadConfigError, config from zarr.core.dtype import data_type_registry -from zarr.errors import ZarrUserWarning +from zarr.errors import UnknownCodecError, ZarrUserWarning if TYPE_CHECKING: from importlib.metadata import EntryPoint @@ -23,7 +24,7 @@ from zarr.abc.numcodec import Numcodec from zarr.core.buffer import Buffer, NDBuffer from zarr.core.chunk_key_encodings import ChunkKeyEncoding - from zarr.core.common import JSON + from zarr.core.common import JSON, ZarrFormat __all__ = [ "Registry", @@ -39,6 +40,131 @@ "register_pipeline", ] +_ZARR_CODEC_DOCS_URL = "https://zarr.readthedocs.io/en/stable/user-guide/extending/#custom-codecs" +_NUMCODECS_CODEC_DOCS_URL = ( + "https://numcodecs.readthedocs.io/en/stable/registry.html#numcodecs.registry.register_codec" +) + +# Codecs zarr-python does not implement, mapped to the names of Python packages that do. +# These tables exist purely to make the "no implementation for this codec" error actionable; +# nothing here affects which codecs zarr can actually read or write. Values are what you would +# pass to `pip install`. Only add an entry you have verified against the package's declared +# entry points, and only for a package that is actually published. +# +# The two Zarr formats resolve codecs through different registries, so they get different +# tables: a name can mean one thing as a Zarr format 3 codec name and another as a Zarr +# format 2 codec id. `imagecodecs_*` is exactly that -- `virtual-tiff` declares 15 of those +# names under `zarr.codecs`, while `imagecodecs-numcodecs` declares all 81 under +# `numcodecs.codecs`, so the format 2 side can use a prefix and the format 3 side cannot. + +# Zarr format 3 codec names (entry point group "zarr.codecs"). +_CODEC_PACKAGES: dict[str, tuple[str, ...]] = { + "gribberish": ("gribberish",), + # `virtual-tiff` declares these 15 `imagecodecs_*` names, out of the 81 that exist as + # numcodecs ids. They are listed exactly rather than by prefix so that the other 66 get no + # hint instead of a hint pointing at a package that does not provide them. + "imagecodecs_deflate": ("virtual-tiff",), + "imagecodecs_delta": ("virtual-tiff",), + "imagecodecs_floatpred": ("virtual-tiff",), + "imagecodecs_jetraw": ("virtual-tiff",), + "imagecodecs_jpeg": ("virtual-tiff",), + "imagecodecs_jpeg2k": ("virtual-tiff",), + "imagecodecs_jpeg8": ("virtual-tiff",), + "imagecodecs_jpegxl": ("virtual-tiff",), + "imagecodecs_jpegxr": ("virtual-tiff",), + "imagecodecs_lerc": ("virtual-tiff",), + "imagecodecs_lzw": ("virtual-tiff",), + "imagecodecs_packbits": ("virtual-tiff",), + "imagecodecs_png": ("virtual-tiff",), + "imagecodecs_webp": ("virtual-tiff",), + "imagecodecs_zstd": ("virtual-tiff",), + "n5_default": ("zarr-n5",), +} + +# As `_CODEC_PACKAGES`, but each key is matched against the start of the codec name. Packages +# that provide many codecs namespace them behind a shared prefix, so one entry covers them all. +_CODEC_PACKAGE_PREFIXES: dict[str, tuple[str, ...]] = { + "any-numcodecs.": ("zarr-any-numcodecs",), + "omfiles.": ("omfiles",), + "virtual_tiff.": ("virtual-tiff",), +} + +# Zarr format 2 codec ids (entry point group "numcodecs.codecs"). `numcodecs` itself gates +# several of its own codecs behind optional dependencies, so the package to install for those +# is an extra of numcodecs rather than a third-party distribution. +_NUMCODEC_PACKAGES: dict[str, tuple[str, ...]] = { + "FITSAscii": ("kerchunk",), + "FITSVarBintable": ("kerchunk",), + "crc32c": ("numcodecs[crc32c]",), + "fill_hdf_strings": ("kerchunk",), + "grib": ("kerchunk",), + "msgpack2": ("numcodecs[msgpack]",), + "pcodec": ("numcodecs[pcodec]",), + "rawgrib": ("gribscan",), + "record_member": ("kerchunk",), + "vc-delta3d": ("vc-delta3d",), + "wavpack": ("wavpack-numcodecs",), + "zfpy": ("numcodecs[zfpy]",), +} + +# As `_NUMCODEC_PACKAGES`, but matched against the start of the codec id. +_NUMCODEC_PACKAGE_PREFIXES: dict[str, tuple[str, ...]] = { + "gribscan.": ("gribscan",), + "imagecodecs_": ("imagecodecs-numcodecs",), +} + + +def _packages_for_codec(name: str, *, zarr_format: ZarrFormat) -> tuple[str, ...]: + """ + Names of Python packages known to provide an implementation of the codec ``name``. + + Returns an empty tuple if we don't know of any. + + Parameters + ---------- + name : str + The codec name (Zarr format 3) or codec id (Zarr format 2) we failed to resolve. + zarr_format : ZarrFormat + Which registry the codec was looked up in. + """ + if zarr_format == 2: + exact, prefixes = _NUMCODEC_PACKAGES, _NUMCODEC_PACKAGE_PREFIXES + else: + exact, prefixes = _CODEC_PACKAGES, _CODEC_PACKAGE_PREFIXES + if name in exact: + return exact[name] + for prefix, packages in prefixes.items(): + if name.startswith(prefix): + return packages + return () + + +def _missing_codec_message(name: str, *, zarr_format: ZarrFormat) -> str: + """ + Build the error message raised when no implementation of the codec ``name`` is available. + + Parameters + ---------- + name : str + The codec name (Zarr format 3) or codec id (Zarr format 2) we failed to resolve. + zarr_format : ZarrFormat + Which registry the codec was looked up in. Zarr format 2 codecs are resolved through + numcodecs, so that case points at the numcodecs registry rather than at zarr's. + """ + if zarr_format == 2: + docs_url, registry = _NUMCODECS_CODEC_DOCS_URL, "numcodecs" + else: + docs_url, registry = _ZARR_CODEC_DOCS_URL, "zarr" + msg = ( + f"An implementation for codec {name!r} is not available. Register one explicitly " + f"using the codec registry (see {docs_url}), or install a Python package that " + f"registers a codec implementation with {registry}." + ) + packages = _packages_for_codec(name, zarr_format=zarr_format) + if packages: + msg += f" Known packages supporting this codec: {', '.join(packages)}." + return msg + class Registry[T](dict[str, type[T]]): def __init__(self) -> None: @@ -168,7 +294,7 @@ def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]: codec_classes = _codec_registries[key] if not codec_classes: - raise KeyError(key) + raise UnknownCodecError(_missing_codec_message(key, zarr_format=3)) config_entry = config.get("codecs", {}).get(key) if config_entry is None: if len(codec_classes) == 1: @@ -179,11 +305,17 @@ def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]: category=ZarrUserWarning, ) return list(codec_classes.values())[-1] - selected_codec_cls = codec_classes[config_entry] - - if selected_codec_cls: - return selected_codec_cls - raise KeyError(key) + selected_codec_cls = codec_classes.get(config_entry) + if selected_codec_cls is None: + # Not UnknownCodecError: the codec is known, the implementation named in the config is + # not registered. That is a configuration problem, which is what the sibling getters in + # this module raise BadConfigError for. + raise BadConfigError( + f"Codec {key!r} is configured to use the implementation {config_entry!r}, which is " + f"not registered. Registered implementations of this codec: " + f"{sorted(codec_classes)}." + ) + return selected_codec_cls def _resolve_codec(data: dict[str, JSON]) -> Codec: @@ -321,6 +453,13 @@ def get_numcodec(data: CodecJSON_V2[str]) -> Numcodec: ------- codec : Numcodec + Raises + ------ + UnknownCodecError + If ``data`` carries a string ``"id"`` that is not registered with numcodecs. Any other + failure, including a registered codec rejecting its configuration and a ``data`` that is + not a mapping, propagates from numcodecs unchanged. + Examples -------- ```python @@ -331,6 +470,19 @@ def get_numcodec(data: CodecJSON_V2[str]) -> Numcodec: ``` """ - from numcodecs.registry import get_codec - + from numcodecs.registry import codec_registry, entries, get_codec + + # Check whether numcodecs can resolve the id *before* handing off, rather than catching what + # `get_codec` raises. Catching cannot tell "this id is unregistered" from "a registered codec + # rejected its configuration" or from "a wrapper codec failed to resolve an inner codec", and + # relabelling either of those with this id would attach a package hint that is simply wrong. + # This mirrors the two lookups `get_codec` performs (it then tests the result for + # truthiness rather than membership, which only differs for a falsy registry value). + # Widened to `object` deliberately: `data` is annotated as a TypedDict, but this is a public + # function and callers pass whatever they like. numcodecs coerces with `dict(config)` and + # raises for anything that is not a mapping, which is the behaviour to preserve. + raw: object = data + codec_id = raw.get("id") if isinstance(raw, Mapping) else None + if isinstance(codec_id, str) and codec_id not in codec_registry and codec_id not in entries: + raise UnknownCodecError(_missing_codec_message(codec_id, zarr_format=2)) return get_codec(data) # type: ignore[no-any-return] diff --git a/tests/test_codecs/test_numcodecs.py b/tests/test_codecs/test_numcodecs.py index 99cd89492f..d78b73f34c 100644 --- a/tests/test_codecs/test_numcodecs.py +++ b/tests/test_codecs/test_numcodecs.py @@ -8,15 +8,10 @@ import pytest from numcodecs import GZip -try: - from numcodecs.errors import UnknownCodecError -except ImportError: - # Older versions of numcodecs don't have a separate errors module - UnknownCodecError = ValueError - from zarr import config, create_array, open_array from zarr.abc.numcodec import _is_numcodec, _is_numcodec_cls from zarr.codecs import numcodecs as _numcodecs +from zarr.errors import UnknownCodecError from zarr.registry import get_codec_class, get_numcodec if TYPE_CHECKING: @@ -297,12 +292,7 @@ def test_generic_checksum(codec_class: type[_numcodecs._NumcodecsBytesBytesCodec def test_generic_bytes_codec(codec_class: type[_numcodecs._NumcodecsArrayBytesCodec]) -> None: try: codec_class()._codec # noqa: B018 - except ValueError as e: # pragma: no cover - if "codec not available" in str(e): - pytest.xfail(f"{codec_class.codec_name} is not available: {e}") # type: ignore[misc] - else: - raise - except ImportError as e: # pragma: no cover + except (UnknownCodecError, ImportError) as e: # pragma: no cover pytest.xfail(f"{codec_class.codec_name} is not available: {e}") # type: ignore[misc] data = np.arange(0, 256, dtype="float32").reshape((16, 16)) diff --git a/tests/test_config.py b/tests/test_config.py index 47f71a798e..c22aac7603 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -24,7 +24,7 @@ from zarr.core.codec_pipeline import BatchedCodecPipeline from zarr.core.config import BadConfigError, config from zarr.core.indexing import SelectorTuple -from zarr.errors import ChunkNotFoundError, ZarrUserWarning +from zarr.errors import ChunkNotFoundError, UnknownCodecError, ZarrUserWarning from zarr.registry import ( fully_qualified_name, get_buffer_class, @@ -334,7 +334,7 @@ class NewCodec2(BytesCodec): pass # error if codec is not registered - with pytest.raises(KeyError): + with pytest.raises(UnknownCodecError): get_codec_class("missing_codec") # no warning if only one implementation is available diff --git a/tests/test_registry.py b/tests/test_registry.py new file mode 100644 index 0000000000..767d3e0439 --- /dev/null +++ b/tests/test_registry.py @@ -0,0 +1,374 @@ +from __future__ import annotations + +import json + +import numcodecs.registry +import pytest + +import zarr +import zarr.registry +from zarr.core.buffer import default_buffer_prototype +from zarr.core.config import BadConfigError, config +from zarr.core.metadata.v3 import parse_codecs +from zarr.errors import UnknownCodecError +from zarr.registry import ( + _CODEC_PACKAGE_PREFIXES, + _CODEC_PACKAGES, + _NUMCODEC_PACKAGE_PREFIXES, + _NUMCODECS_CODEC_DOCS_URL, + _ZARR_CODEC_DOCS_URL, + _missing_codec_message, + _packages_for_codec, + get_codec_class, + get_numcodec, +) +from zarr.storage import MemoryStore + + +@pytest.fixture +def unregistered_v3_codec(monkeypatch: pytest.MonkeyPatch) -> str: + """Guarantee a Zarr format 3 codec name resolves to nothing. + + The registry is entry-point driven, so a name the hint table advertises resolves for real + in any environment where the advertised package happens to be installed. + """ + name = "n5_default" + monkeypatch.setitem(zarr.registry._codec_registries, name, zarr.registry.Registry()) + return name + + +@pytest.fixture +def unregistered_v2_codec(monkeypatch: pytest.MonkeyPatch) -> str: + """As above, for a numcodecs codec id.""" + codec_id = "wavpack" + monkeypatch.delitem(numcodecs.registry.codec_registry, codec_id, raising=False) + monkeypatch.delitem(numcodecs.registry.entries, codec_id, raising=False) + return codec_id + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("n5_default", ("zarr-n5",)), + ("gribberish", ("gribberish",)), + ("imagecodecs_jpeg2k", ("virtual-tiff",)), + ("omfiles.pfor", ("omfiles",)), + ("any-numcodecs.array-array", ("zarr-any-numcodecs",)), + ("totally-made-up", ()), + ], +) +def test_packages_for_codec_v3(name: str, expected: tuple[str, ...]) -> None: + """Exact names and prefixes both resolve; unknown names resolve to nothing.""" + assert _packages_for_codec(name, zarr_format=3) == expected + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("wavpack", ("wavpack-numcodecs",)), + ("grib", ("kerchunk",)), + ("zfpy", ("numcodecs[zfpy]",)), + ("crc32c", ("numcodecs[crc32c]",)), + ("gribscan.rawgrib", ("gribscan",)), + ("imagecodecs_jpeg2k", ("imagecodecs-numcodecs",)), + ("totally-made-up", ()), + ], +) +def test_packages_for_numcodec_v2(name: str, expected: tuple[str, ...]) -> None: + """Zarr format 2 codec ids resolve against the numcodecs table.""" + assert _packages_for_codec(name, zarr_format=2) == expected + + +def test_packages_for_codec_is_format_specific() -> None: + """The same name can mean different packages in each format's registry.""" + assert _packages_for_codec("imagecodecs_jpeg2k", zarr_format=3) == ("virtual-tiff",) + assert _packages_for_codec("imagecodecs_jpeg2k", zarr_format=2) == ("imagecodecs-numcodecs",) + # `crc32c` is a codec zarr implements in format 3, so only format 2 gets a hint for it. + assert _packages_for_codec("crc32c", zarr_format=3) == () + + +def test_missing_codec_message_with_known_packages() -> None: + """The message names the codec, the docs page, and every known package.""" + msg = _missing_codec_message("n5_default", zarr_format=3) + assert msg == ( + "An implementation for codec 'n5_default' is not available. Register one explicitly " + f"using the codec registry (see {_ZARR_CODEC_DOCS_URL}), or install a Python package " + "that registers a codec implementation with zarr. Known packages supporting this " + "codec: zarr-n5." + ) + + +def test_missing_codec_message_without_known_packages() -> None: + """With no known package we still explain how to register one by hand.""" + msg = _missing_codec_message("totally-made-up", zarr_format=3) + assert msg == ( + "An implementation for codec 'totally-made-up' is not available. Register one " + f"explicitly using the codec registry (see {_ZARR_CODEC_DOCS_URL}), or install a " + "Python package that registers a codec implementation with zarr." + ) + assert "Known packages" not in msg + + +def test_missing_codec_message_for_zarr_format_2() -> None: + """Format 2 codecs live in the numcodecs registry, so the message must say so.""" + msg = _missing_codec_message("wavpack", zarr_format=2) + assert msg == ( + "An implementation for codec 'wavpack' is not available. Register one explicitly " + f"using the codec registry (see {_NUMCODECS_CODEC_DOCS_URL}), or install a Python " + "package that registers a codec implementation with numcodecs. Known packages " + "supporting this codec: wavpack-numcodecs." + ) + + +def test_no_prefix_shadows_another_prefix() -> None: + """First-match prefix lookup is only deterministic while no prefix contains another.""" + for table in (_CODEC_PACKAGE_PREFIXES, _NUMCODEC_PACKAGE_PREFIXES): + for a in table: + for b in table: + assert a == b or not a.startswith(b) + + +def test_mapping_does_not_shadow_builtin_codecs() -> None: + """A codec zarr implements itself must never appear in the format 3 hint table.""" + import zarr.codecs # noqa: F401 (importing registers the built-in codecs) + from zarr.registry import _codec_registries + + # Select on the implementing class's module rather than on "is this registry non-empty": + # a third-party entry-point codec lazy-loaded by an earlier test would otherwise show up + # here and fail the assertion even though it shadows nothing. + implemented = { + name + for name, reg in _codec_registries.items() + if any(cls.__module__.startswith("zarr.") for cls in reg.values()) + } + assert implemented, "expected importing zarr.codecs to populate the registry" + assert not (implemented & set(_CODEC_PACKAGES)) + for prefix in _CODEC_PACKAGE_PREFIXES: + assert not any(name.startswith(prefix) for name in implemented) + + +def test_get_codec_class_unknown_raises_with_package_hint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An unregistered codec with a known package names that package in the error.""" + from collections import defaultdict + + monkeypatch.setattr(zarr.registry, "_codec_registries", defaultdict(zarr.registry.Registry)) + with pytest.raises(UnknownCodecError, match="Known packages supporting this codec: zarr-n5"): + get_codec_class("n5_default") + + +def test_get_codec_class_unknown_raises_without_package_hint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An unregistered codec we know nothing about still explains manual registration.""" + from collections import defaultdict + + monkeypatch.setattr(zarr.registry, "_codec_registries", defaultdict(zarr.registry.Registry)) + with pytest.raises(UnknownCodecError, match="An implementation for codec 'nope' is not"): + get_codec_class("nope") + + +def test_unknown_codec_error_is_exported() -> None: + """UnknownCodecError is public API now that we tell users to catch it.""" + import zarr.errors + + assert "UnknownCodecError" in zarr.errors.__all__ + + +def test_get_numcodec_unknown_raises_with_package_hint(unregistered_v2_codec: str) -> None: + """A Zarr format 2 codec id we don't have names the package that provides it.""" + with pytest.raises( + UnknownCodecError, match="Known packages supporting this codec: wavpack-numcodecs" + ): + get_numcodec({"id": unregistered_v2_codec}) + + +def test_get_numcodec_unknown_points_at_numcodecs_registry() -> None: + """The Zarr format 2 message links the numcodecs registry, not zarr's extending guide.""" + with pytest.raises(UnknownCodecError) as excinfo: + get_numcodec({"id": "definitely-not-a-real-codec"}) + assert _NUMCODECS_CODEC_DOCS_URL in str(excinfo.value) + assert _ZARR_CODEC_DOCS_URL not in str(excinfo.value) + + +def test_get_numcodec_known_codec_still_works() -> None: + """The happy path is untouched.""" + from numcodecs import GZip + + assert get_numcodec({"id": "gzip", "level": 2}) == GZip(level=2) # type: ignore[typeddict-unknown-key] + + +def test_get_numcodec_without_an_id_keeps_the_numcodecs_error() -> None: + """With no codec id there is nothing to look up, so numcodecs' own error stands.""" + # Not asserting on numcodecs' exception class: it only gained a dedicated + # `numcodecs.errors.UnknownCodecError` in 0.15.1, and zarr supports numcodecs >= 0.14. + with pytest.raises(ValueError) as excinfo: + get_numcodec({"level": 2}) # type: ignore[typeddict-item,typeddict-unknown-key] + assert not isinstance(excinfo.value, UnknownCodecError) + + +def test_get_numcodec_does_not_relabel_a_bad_configuration() -> None: + """A registered codec rejecting its config is not a missing codec. + + Telling the user to install a package they already have would be the same misleading + error this module exists to remove, pointing the other way. + """ + with pytest.raises(ValueError) as excinfo: + get_numcodec({"id": "bitround", "keepbits": -1}) # type: ignore[typeddict-unknown-key] + assert not isinstance(excinfo.value, UnknownCodecError) + assert str(excinfo.value) == "keepbits must be zero or positive" + + +@pytest.mark.parametrize("via", ["zarr", "numcodecs"]) +def test_get_numcodec_does_not_relabel_a_missing_inner_codec(via: str) -> None: + """A wrapper codec failing on an inner codec must not be blamed on the outer id. + + The outer id is registered, so relabelling it would advertise a package the user has + already installed. Covers both routes a real wrapper takes to resolve its inner codec: + zarr's ``get_numcodec`` and numcodecs' own ``get_codec``. + """ + import numcodecs.registry + from numcodecs.abc import Codec + + resolve = get_numcodec if via == "zarr" else numcodecs.registry.get_codec + + class WrapperCodec(Codec): # type: ignore[misc] + codec_id = "test_wrapper" + + @classmethod + def from_config(cls, config: dict[str, object]) -> WrapperCodec: + resolve({"id": "some-missing-inner-codec"}) + return cls() + + def encode(self, buf: object) -> object: + return buf + + def decode(self, buf: object, out: object = None) -> object: + return buf + + numcodecs.registry.register_codec(WrapperCodec, codec_id="test_wrapper") + try: + with pytest.raises(ValueError) as excinfo: + get_numcodec({"id": "test_wrapper"}) + assert "some-missing-inner-codec" in str(excinfo.value) + assert "test_wrapper" not in str(excinfo.value) + finally: + numcodecs.registry.codec_registry.pop("test_wrapper", None) + + +@pytest.mark.parametrize("data", ["abc", ["ab"]]) +def test_get_numcodec_non_mapping_input_still_raises_value_error(data: object) -> None: + """Input that is not a mapping must reach numcodecs rather than raising AttributeError. + + Reading ``id`` off the input unconditionally turned numcodecs' ValueError into an + AttributeError, which is not a ValueError and so escaped handlers that caught it. Both + params end in a ValueError from numcodecs, by different routes: ``"abc"`` fails + ``dict(config)`` coercion, while ``["ab"]`` coerces to ``{"a": "b"}`` and then has no id. + """ + with pytest.raises(ValueError): + get_numcodec(data) # type: ignore[arg-type] + + +def test_imagecodecs_prefix_does_not_over_match_in_zarr_format_3() -> None: + """virtual-tiff provides 15 of the 81 `imagecodecs_*` names; the rest must get no hint. + + Recommending virtual-tiff for a name it does not provide is worse than saying nothing. + """ + assert _packages_for_codec("imagecodecs_jpeg2k", zarr_format=3) == ("virtual-tiff",) + for name in ("imagecodecs_jpegls", "imagecodecs_avif", "imagecodecs_blosc"): + assert _packages_for_codec(name, zarr_format=3) == () + assert _packages_for_codec(name, zarr_format=2) == ("imagecodecs-numcodecs",) + + +def test_resolve_codec_reports_missing_codec() -> None: + """The other public entry point into `get_codec_class` gets the message too.""" + import zarr + + with pytest.raises(UnknownCodecError, match="Known packages supporting this codec: zarr-n5"): + zarr.create_array({}, shape=(4,), dtype="uint8", compressors=[{"name": "n5_default"}]) + + +async def test_open_array_with_missing_v3_codec_reports_package( + unregistered_v3_codec: str, +) -> None: + """Opening a Zarr format 3 array naming a codec we lack points at the package.""" + store = MemoryStore() + metadata = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "data_type": "float64", + "fill_value": 0.0, + "codecs": [{"name": "bytes"}, {"name": unregistered_v3_codec}], + "attributes": {}, + } + await store.set( + "zarr.json", + default_buffer_prototype().buffer.from_bytes(json.dumps(metadata).encode()), + ) + with pytest.raises(UnknownCodecError, match="Known packages supporting this codec: zarr-n5"): + zarr.open_array(store=store, mode="r") + + +async def test_open_array_with_missing_v2_codec_reports_package( + unregistered_v2_codec: str, +) -> None: + """Same, for a Zarr format 2 array whose compressor id we lack.""" + store = MemoryStore() + metadata = { + "zarr_format": 2, + "shape": [4], + "chunks": [4], + "dtype": " None: + """A config pinning an implementation that isn't registered is a config error.""" + with config.set({"codecs.bytes": "some.package.RemovedBytesCodec"}): + with pytest.raises(BadConfigError, match="RemovedBytesCodec"): + get_codec_class("bytes") + + +def test_parse_codecs_with_unregistered_config_pin_raises_bad_config_error() -> None: + """The same, through the array-open path, and never as a bare KeyError.""" + with config.set({"codecs.bytes": "some.package.RemovedBytesCodec"}): + with pytest.raises(BadConfigError, match="RemovedBytesCodec"): + parse_codecs([{"name": "bytes"}]) + + +def test_parse_codecs_converts_keyerror_from_from_dict(monkeypatch: pytest.MonkeyPatch) -> None: + """A codec whose from_dict indexes a malformed config must not leak a KeyError. + + A bare KeyError out of metadata parsing is caught by the array-then-group fallback in + `zarr.api.asynchronous.open`, which then reports an unrelated group error. + """ + from zarr.codecs import BytesCodec + from zarr.errors import MetadataValidationError + from zarr.registry import register_codec + + class PickyCodec(BytesCodec): + @classmethod + def from_dict(cls, data: object) -> PickyCodec: + data["configuration"]["required_option"] # type: ignore[index] + return cls() + + monkeypatch.setitem(zarr.registry._codec_registries, "test_picky", zarr.registry.Registry()) + register_codec("test_picky", PickyCodec) + with pytest.raises(MetadataValidationError, match="test_picky.*required_option"): + parse_codecs([{"name": "test_picky", "configuration": {}}]) From 5e00e8ccd7269314eb23fb21ae50b8ec7847f6fd Mon Sep 17 00:00:00 2001 From: Theodoros Visvikis <129113023+Tomatokeftes@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:07:06 +0200 Subject: [PATCH 3/5] Fix infinite loop in _guess_num_chunks_per_axis_shard for 0-dimensional arrays (#4305) * Fix infinite loop in shard size guessing for 0-dimensional arrays _guess_num_chunks_per_axis_shard looped forever for rank-0 arrays because neither loop condition depends on the counter when chunk_shape is empty. Return 1 immediately in that case, which makes the budget path produce shards=() just like the no-budget path. Closes #4304 * Add changelog fragment --- changes/4305.bugfix.md | 1 + src/zarr/core/chunk_grids.py | 2 ++ tests/test_chunk_grids.py | 26 ++++++++++++++++++++++++++ 3 files changed, 29 insertions(+) create mode 100644 changes/4305.bugfix.md diff --git a/changes/4305.bugfix.md b/changes/4305.bugfix.md new file mode 100644 index 0000000000..13221d8e3b --- /dev/null +++ b/changes/4305.bugfix.md @@ -0,0 +1 @@ +Fixed an infinite loop when creating a 0-dimensional array with `shards="auto"` while the `array.target_shard_size_bytes` config option is set. Such arrays now resolve to `shards=()`, matching the behavior when no shard size target is configured. diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 584829bc6c..5e19364cae 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -865,6 +865,8 @@ def _guess_num_chunks_per_axis_shard( if max_bytes < bytes_per_chunk: return 1 num_axes = len(chunk_shape) + if num_axes == 0: + return 1 chunks_per_shard = 1 # First check for byte size, second check to make sure we don't go bigger than the array shape while (bytes_per_chunk * ((chunks_per_shard + 1) ** num_axes)) <= max_bytes and all( diff --git a/tests/test_chunk_grids.py b/tests/test_chunk_grids.py index 4640c43d1c..4c3c437eb2 100644 --- a/tests/test_chunk_grids.py +++ b/tests/test_chunk_grids.py @@ -3,14 +3,17 @@ import numpy as np import pytest +import zarr from tests.conftest import Expect, ExpectFail from zarr.core.chunk_grids import ( ChunkLayout, + _guess_num_chunks_per_axis_shard, _guess_regular_chunks, normalize_chunks_1d, normalize_chunks_nd, resolve_outer_and_inner_chunks, ) +from zarr.errors import ZarrUserWarning def _assert_chunks_equal( @@ -277,3 +280,26 @@ def test_normalize_chunks_1d_returns_int64_array( assert result.dtype == np.int64 assert result.ndim == 1 assert result.tolist() == case.output + + +def test_guess_num_chunks_per_axis_shard_0d() -> None: + """Regression test for https://github.com/zarr-developers/zarr-python/issues/4304.""" + assert ( + _guess_num_chunks_per_axis_shard( + chunk_shape=(), item_size=8, max_bytes=128 * 1024 * 1024, array_shape=() + ) + == 1 + ) + + +def test_create_0d_array_auto_shards_with_target_shard_size() -> None: + """A 0-dimensional array with shards="auto" and a shard size budget must not hang. + + Regression test for https://github.com/zarr-developers/zarr-python/issues/4304. + """ + with ( + zarr.config.set({"array.target_shard_size_bytes": 128 * 1024 * 1024}), + pytest.warns(ZarrUserWarning, match="Automatic shard shape inference is experimental"), + ): + arr = zarr.create_array(store={}, shape=(), dtype="int64", shards="auto") + assert arr.shards == () From 9058ec8b0f603833f6358d7ba286ffe926fcc163 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 2 Sep 2026 13:11:47 +0200 Subject: [PATCH 4/5] fix: resolve chunks=-1 to chunk size 1 on zero-length axes The -1 chunk sentinel resolved to chunk size 0 on zero-length axes, which broke every downstream sharding path differently: a ValueError with shards=None, a ZeroDivisionError with shards="auto", and an infinite loop with shards="auto" plus array.target_shard_size_bytes. Clamp it to 1, matching the auto-chunking clamp in _guess_regular_chunks. Also guard _guess_num_chunks_per_axis_shard against zero-size chunk shapes directly (same non-terminating-loop cause as the 0-d case fixed in #4305), and fix the arithmetic in its docstring example. Follow-up to #4305 / issue #4304. Assisted-by: ClaudeCode:claude-fable-5-1 --- src/zarr/core/chunk_grids.py | 13 ++++++++--- tests/test_chunk_grids.py | 44 +++++++++++++++++++++++++++++++++--- 2 files changed, 51 insertions(+), 6 deletions(-) diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 5e19364cae..13ec3d8ffa 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -734,7 +734,9 @@ def normalize_chunks_1d( if isinstance(chunks, numbers.Integral): chunk_size = int(chunks) if chunk_size == -1: - return np.array([span], dtype=np.int64) + # A zero-length span still gets chunk size 1 (chunk sizes must be positive), + # matching the auto-chunking clamp in _guess_regular_chunks. + return np.array([max(span, 1)], dtype=np.int64) if chunk_size <= 0: raise ValueError(f"Chunk size must be positive, got {chunk_size}") if span == 0: @@ -844,7 +846,10 @@ def _guess_num_chunks_per_axis_shard( For example, for a (2,2,2) chunk size and item size 4, maximum bytes of 256 would return 2. In other words the shard would be a (2,2,2) grid of (2,2,2) chunks - i.e., prod(chunk_shape) * (returned_val * len(chunk_shape)) * item_size = 256 bytes. + i.e., prod(chunk_shape) * (returned_val ** len(chunk_shape)) * item_size = 256 bytes. + + Degenerate chunk shapes — a 0-dimensional shape, or one containing a zero-length + axis — return 1, as the search loop's stopping conditions can never be met. Parameters ---------- @@ -865,7 +870,9 @@ def _guess_num_chunks_per_axis_shard( if max_bytes < bytes_per_chunk: return 1 num_axes = len(chunk_shape) - if num_axes == 0: + # For a 0-dimensional chunk shape or one with a zero-length axis, both loop + # conditions below are constant, so the loop would never terminate. + if num_axes == 0 or bytes_per_chunk == 0: return 1 chunks_per_shard = 1 # First check for byte size, second check to make sure we don't go bigger than the array shape diff --git a/tests/test_chunk_grids.py b/tests/test_chunk_grids.py index 4c3c437eb2..0f1309cf05 100644 --- a/tests/test_chunk_grids.py +++ b/tests/test_chunk_grids.py @@ -267,6 +267,8 @@ def test_normalize_chunks_nd_errors(case: ExpectFail[tuple[Any, tuple[int, ...]] Expect(input=([10, 20, 30, 40], 100), output=[10, 20, 30, 40], id="explicit-list"), # -1 sentinel branch: one chunk covering the full span. Expect(input=(-1, 100), output=[100], id="full-span-sentinel"), + # -1 on a zero-length span clamps to chunk size 1 (chunk sizes must be positive). + Expect(input=(-1, 0), output=[1], id="full-span-sentinel-empty"), ], ids=lambda c: c.id, ) @@ -282,11 +284,24 @@ def test_normalize_chunks_1d_returns_int64_array( assert result.tolist() == case.output -def test_guess_num_chunks_per_axis_shard_0d() -> None: - """Regression test for https://github.com/zarr-developers/zarr-python/issues/4304.""" +@pytest.mark.parametrize( + ("chunk_shape", "array_shape"), + [((), ()), ((0,), (0,)), ((0, 0), (0, 0))], + ids=["0d", "zero-1d", "zero-2d"], +) +def test_guess_num_chunks_per_axis_shard_degenerate( + chunk_shape: tuple[int, ...], array_shape: tuple[int, ...] +) -> None: + """Degenerate chunk shapes must return 1 instead of hanging the search loop. + + Regression test for https://github.com/zarr-developers/zarr-python/issues/4304. + """ assert ( _guess_num_chunks_per_axis_shard( - chunk_shape=(), item_size=8, max_bytes=128 * 1024 * 1024, array_shape=() + chunk_shape=chunk_shape, + item_size=8, + max_bytes=128 * 1024 * 1024, + array_shape=array_shape, ) == 1 ) @@ -303,3 +318,26 @@ def test_create_0d_array_auto_shards_with_target_shard_size() -> None: ): arr = zarr.create_array(store={}, shape=(), dtype="int64", shards="auto") assert arr.shards == () + + +@pytest.mark.parametrize( + "target_shard_size_bytes", + [None, 128 * 1024 * 1024], + ids=["no-budget", "budget"], +) +def test_create_zero_length_array_full_span_chunks_auto_shards( + target_shard_size_bytes: int | None, +) -> None: + """`chunks=-1` on a zero-length axis with shards="auto" must neither hang nor raise. + + The -1 sentinel used to resolve to chunk size 0 on zero-length axes, which broke + every sharding code path: a ZeroDivisionError without a shard size budget, and an + infinite loop with one (https://github.com/zarr-developers/zarr-python/issues/4304). + """ + with ( + zarr.config.set({"array.target_shard_size_bytes": target_shard_size_bytes}), + pytest.warns(ZarrUserWarning, match="Automatic shard shape inference is experimental"), + ): + arr = zarr.create_array(store={}, shape=(0,), dtype="int64", chunks=-1, shards="auto") + assert arr.chunks == (1,) + assert arr.shards == (1,) From f850df574a208642cfee61a0612cd8e7b6e5604e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 2 Sep 2026 13:12:14 +0200 Subject: [PATCH 5/5] docs: add changelog fragment Assisted-by: ClaudeCode:claude-fable-5-1 --- changes/4307.bugfix.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/4307.bugfix.md diff --git a/changes/4307.bugfix.md b/changes/4307.bugfix.md new file mode 100644 index 0000000000..b77c89cfa5 --- /dev/null +++ b/changes/4307.bugfix.md @@ -0,0 +1 @@ +Fixed `chunks=-1` on a zero-length axis resolving to an invalid chunk size of 0, which caused a `ValueError`, `ZeroDivisionError`, or infinite loop depending on the sharding configuration. Such axes now get chunk size 1, matching `chunks="auto"`.