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/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/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"`. 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/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/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 584829bc6c..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,6 +870,10 @@ def _guess_num_chunks_per_axis_shard( if max_bytes < bytes_per_chunk: return 1 num_axes = len(chunk_shape) + # 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 while (bytes_per_chunk * ((chunks_per_shard + 1) ** num_axes)) <= max_bytes and all( 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_chunk_grids.py b/tests/test_chunk_grids.py index 4640c43d1c..0f1309cf05 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( @@ -264,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, ) @@ -277,3 +282,62 @@ def test_normalize_chunks_1d_returns_int64_array( assert result.dtype == np.int64 assert result.ndim == 1 assert result.tolist() == case.output + + +@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=chunk_shape, + item_size=8, + max_bytes=128 * 1024 * 1024, + array_shape=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 == () + + +@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,) 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_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 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": {}}])