From 794b5a2d370d0901ddf2f7e3be8edf6965ff5821 Mon Sep 17 00:00:00 2001 From: victorwon2001 <192616110+victorwon2001@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:34:42 +0900 Subject: [PATCH] refactor: centralize built-in codec registration --- src/zarr/__init__.py | 3 + src/zarr/codecs/__init__.py | 50 +-------------- src/zarr/registry.py | 66 ++++++++++++++++++++ tests/test_registry.py | 118 ++++++++++++++++++++++++++++++++++++ 4 files changed, 190 insertions(+), 47 deletions(-) diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py index f305245f4e..975a8d7416 100644 --- a/src/zarr/__init__.py +++ b/src/zarr/__init__.py @@ -37,6 +37,9 @@ from zarr.core.array import Array, AsyncArray from zarr.core.config import config from zarr.core.group import AsyncGroup, Group +from zarr.registry import _register_builtin_codecs + +_register_builtin_codecs() # in case setuptools scm screw up and find version to be 0.0.0 if __version__.startswith("0.0.0"): diff --git a/src/zarr/codecs/__init__.py b/src/zarr/codecs/__init__.py index 9a1b47b351..b78a9092d4 100644 --- a/src/zarr/codecs/__init__.py +++ b/src/zarr/codecs/__init__.py @@ -5,7 +5,9 @@ from zarr.codecs.cast_value import CastValue from zarr.codecs.crc32c_ import Crc32cCodec from zarr.codecs.gzip import GzipCodec -from zarr.codecs.numcodecs import ( + +# Keep the numcodecs wrappers importable from this module for compatibility. +from zarr.codecs.numcodecs import ( # noqa: F401 BZ2, CRC32, CRC32C, @@ -33,7 +35,6 @@ from zarr.codecs.transpose import TransposeCodec from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec from zarr.codecs.zstd import ZstdCodec -from zarr.registry import register_codec __all__ = [ "BloscCname", @@ -53,48 +54,3 @@ "VLenUTF8Codec", "ZstdCodec", ] - -register_codec("blosc", BloscCodec) -register_codec("cast_value", CastValue) -register_codec("bytes", BytesCodec) - -# compatibility with earlier versions of ZEP1 -register_codec("endian", BytesCodec) -register_codec("crc32c", Crc32cCodec) -register_codec("gzip", GzipCodec) -register_codec("scale_offset", ScaleOffset) -register_codec("sharding_indexed", ShardingCodec) -register_codec("zstd", ZstdCodec) -register_codec("vlen-utf8", VLenUTF8Codec) -register_codec("vlen-bytes", VLenBytesCodec) -register_codec("transpose", TransposeCodec) - -# Register all the codecs formerly contained in numcodecs.zarr3 - -register_codec("numcodecs.bz2", BZ2, qualname="zarr.codecs.numcodecs.BZ2") -register_codec("numcodecs.crc32", CRC32, qualname="zarr.codecs.numcodecs.CRC32") -register_codec("numcodecs.crc32c", CRC32C, qualname="zarr.codecs.numcodecs.CRC32C") -register_codec("numcodecs.lz4", LZ4, qualname="zarr.codecs.numcodecs.LZ4") -register_codec("numcodecs.lzma", LZMA, qualname="zarr.codecs.numcodecs.LZMA") -register_codec("numcodecs.zfpy", ZFPY, qualname="zarr.codecs.numcodecs.ZFPY") -register_codec("numcodecs.adler32", Adler32, qualname="zarr.codecs.numcodecs.Adler32") -register_codec("numcodecs.astype", AsType, qualname="zarr.codecs.numcodecs.AsType") -register_codec("numcodecs.bitround", BitRound, qualname="zarr.codecs.numcodecs.BitRound") -register_codec("numcodecs.blosc", Blosc, qualname="zarr.codecs.numcodecs.Blosc") -register_codec("numcodecs.delta", Delta, qualname="zarr.codecs.numcodecs.Delta") -register_codec( - "numcodecs.fixedscaleoffset", - FixedScaleOffset, - qualname="zarr.codecs.numcodecs.FixedScaleOffset", -) -register_codec("numcodecs.fletcher32", Fletcher32, qualname="zarr.codecs.numcodecs.Fletcher32") -register_codec("numcodecs.gzip", GZip, qualname="zarr.codecs.numcodecs.GZip") -register_codec( - "numcodecs.jenkins_lookup3", JenkinsLookup3, qualname="zarr.codecs.numcodecs.JenkinsLookup3" -) -register_codec("numcodecs.pcodec", PCodec, qualname="zarr.codecs.numcodecs.PCodec") -register_codec("numcodecs.packbits", PackBits, qualname="zarr.codecs.numcodecs.PackBits") -register_codec("numcodecs.quantize", Quantize, qualname="zarr.codecs.numcodecs.Quantize") -register_codec("numcodecs.shuffle", Shuffle, qualname="zarr.codecs.numcodecs.Shuffle") -register_codec("numcodecs.zlib", Zlib, qualname="zarr.codecs.numcodecs.Zlib") -register_codec("numcodecs.zstd", Zstd, qualname="zarr.codecs.numcodecs.Zstd") diff --git a/src/zarr/registry.py b/src/zarr/registry.py index 307c323153..ac352cd473 100644 --- a/src/zarr/registry.py +++ b/src/zarr/registry.py @@ -253,6 +253,7 @@ def register(self, cls: type[T], qualname: str | None = None) -> None: _codec_registries: dict[str, Registry[Codec]] = defaultdict(Registry) +_builtin_codecs_registered = False _pipeline_registry: Registry[CodecPipeline] = Registry() _buffer_registry: Registry[Buffer] = Registry() _ndbuffer_registry: Registry[NDBuffer] = Registry() @@ -337,6 +338,71 @@ def register_codec(key: str, codec_cls: type[Codec], *, qualname: str | None = N _codec_registries[key].register(codec_cls, qualname=qualname) +def _register_builtin_codecs() -> None: + """Register built-in codecs after the registry and codec implementations are initialized.""" + global _builtin_codecs_registered + if _builtin_codecs_registered: + return + + from zarr import codecs + from zarr.codecs import numcodecs + + register_codec("blosc", codecs.BloscCodec) + register_codec("cast_value", codecs.CastValue) + register_codec("bytes", codecs.BytesCodec) + + # Compatibility with earlier versions of ZEP1. + register_codec("endian", codecs.BytesCodec) + register_codec("crc32c", codecs.Crc32cCodec) + register_codec("gzip", codecs.GzipCodec) + register_codec("scale_offset", codecs.ScaleOffset) + register_codec("sharding_indexed", codecs.ShardingCodec) + register_codec("zstd", codecs.ZstdCodec) + register_codec("vlen-utf8", codecs.VLenUTF8Codec) + register_codec("vlen-bytes", codecs.VLenBytesCodec) + register_codec("transpose", codecs.TransposeCodec) + + # Compatibility with codecs formerly contained in numcodecs.zarr3. + register_codec("numcodecs.bz2", numcodecs.BZ2, qualname="zarr.codecs.numcodecs.BZ2") + register_codec("numcodecs.crc32", numcodecs.CRC32, qualname="zarr.codecs.numcodecs.CRC32") + register_codec("numcodecs.crc32c", numcodecs.CRC32C, qualname="zarr.codecs.numcodecs.CRC32C") + register_codec("numcodecs.lz4", numcodecs.LZ4, qualname="zarr.codecs.numcodecs.LZ4") + register_codec("numcodecs.lzma", numcodecs.LZMA, qualname="zarr.codecs.numcodecs.LZMA") + register_codec("numcodecs.zfpy", numcodecs.ZFPY, qualname="zarr.codecs.numcodecs.ZFPY") + register_codec("numcodecs.adler32", numcodecs.Adler32, qualname="zarr.codecs.numcodecs.Adler32") + register_codec("numcodecs.astype", numcodecs.AsType, qualname="zarr.codecs.numcodecs.AsType") + register_codec( + "numcodecs.bitround", numcodecs.BitRound, qualname="zarr.codecs.numcodecs.BitRound" + ) + register_codec("numcodecs.blosc", numcodecs.Blosc, qualname="zarr.codecs.numcodecs.Blosc") + register_codec("numcodecs.delta", numcodecs.Delta, qualname="zarr.codecs.numcodecs.Delta") + register_codec( + "numcodecs.fixedscaleoffset", + numcodecs.FixedScaleOffset, + qualname="zarr.codecs.numcodecs.FixedScaleOffset", + ) + register_codec( + "numcodecs.fletcher32", numcodecs.Fletcher32, qualname="zarr.codecs.numcodecs.Fletcher32" + ) + register_codec("numcodecs.gzip", numcodecs.GZip, qualname="zarr.codecs.numcodecs.GZip") + register_codec( + "numcodecs.jenkins_lookup3", + numcodecs.JenkinsLookup3, + qualname="zarr.codecs.numcodecs.JenkinsLookup3", + ) + register_codec("numcodecs.pcodec", numcodecs.PCodec, qualname="zarr.codecs.numcodecs.PCodec") + register_codec( + "numcodecs.packbits", numcodecs.PackBits, qualname="zarr.codecs.numcodecs.PackBits" + ) + register_codec( + "numcodecs.quantize", numcodecs.Quantize, qualname="zarr.codecs.numcodecs.Quantize" + ) + register_codec("numcodecs.shuffle", numcodecs.Shuffle, qualname="zarr.codecs.numcodecs.Shuffle") + register_codec("numcodecs.zlib", numcodecs.Zlib, qualname="zarr.codecs.numcodecs.Zlib") + register_codec("numcodecs.zstd", numcodecs.Zstd, qualname="zarr.codecs.numcodecs.Zstd") + _builtin_codecs_registered = True + + def register_pipeline(pipe_cls: type[CodecPipeline]) -> None: _pipeline_registry.register(pipe_cls) diff --git a/tests/test_registry.py b/tests/test_registry.py index 4344f5d72d..bfd07d3d6a 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -1,6 +1,9 @@ from __future__ import annotations import json +import subprocess +import sys +from importlib.metadata import EntryPoint import numcodecs.registry import pytest @@ -25,6 +28,121 @@ from zarr.storage import MemoryStore +@pytest.mark.parametrize( + "imports", + [ + "import zarr.registry\nimport zarr.codecs", + "import zarr.codecs\nimport zarr.registry", + "import zarr", + ], + ids=["registry-first", "codecs-first", "zarr-only"], +) +def test_builtin_codecs_register_in_any_import_order(imports: str) -> None: + """The package bootstrap registers each codec even when registry is imported first.""" + builtin = { + "blosc": "BloscCodec", + "cast_value": "CastValue", + "bytes": "BytesCodec", + "endian": "BytesCodec", + "crc32c": "Crc32cCodec", + "gzip": "GzipCodec", + "scale_offset": "ScaleOffset", + "sharding_indexed": "ShardingCodec", + "zstd": "ZstdCodec", + "vlen-utf8": "VLenUTF8Codec", + "vlen-bytes": "VLenBytesCodec", + "transpose": "TransposeCodec", + } + numcodecs = { + "bz2": "BZ2", + "crc32": "CRC32", + "crc32c": "CRC32C", + "lz4": "LZ4", + "lzma": "LZMA", + "zfpy": "ZFPY", + "adler32": "Adler32", + "astype": "AsType", + "bitround": "BitRound", + "blosc": "Blosc", + "delta": "Delta", + "fixedscaleoffset": "FixedScaleOffset", + "fletcher32": "Fletcher32", + "gzip": "GZip", + "jenkins_lookup3": "JenkinsLookup3", + "pcodec": "PCodec", + "packbits": "PackBits", + "quantize": "Quantize", + "shuffle": "Shuffle", + "zlib": "Zlib", + "zstd": "Zstd", + } + script = f""" +{imports} +import zarr.codecs as codecs +import zarr.codecs.numcodecs as numcodecs +from zarr.registry import _codec_registries, _register_builtin_codecs, get_codec_class + +for name, class_name in {builtin!r}.items(): + cls = getattr(codecs, class_name) + assert get_codec_class(name) is cls, name + assert _codec_registries[name][f'{{cls.__module__}}.{{cls.__qualname__}}'] is cls, name +for name, class_name in {numcodecs!r}.items(): + cls = getattr(numcodecs, class_name) + key = f'numcodecs.{{name}}' + assert get_codec_class(key) is cls, key + assert _codec_registries[key][f'zarr.codecs.numcodecs.{{class_name}}'] is cls, key +snapshot = {{key: dict(value) for key, value in _codec_registries.items()}} +_register_builtin_codecs() +assert snapshot == {{key: dict(value) for key, value in _codec_registries.items()}} +""" + completed = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, check=False + ) + assert completed.returncode == 0, completed.stderr + + +def test_builtin_codec_bootstrap_preserves_custom_registration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Running bootstrap again must not replace user-provided implementations.""" + from zarr.codecs import BytesCodec + from zarr.registry import Registry, _codec_registries, _register_builtin_codecs, register_codec + + class CustomBytesCodec(BytesCodec): + pass + + monkeypatch.setitem( + _codec_registries["bytes"], "zarr.codecs.bytes.BytesCodec", CustomBytesCodec + ) + monkeypatch.setitem(_codec_registries, "custom_builtin_test", Registry()) + register_codec("custom_builtin_test", CustomBytesCodec) + before = {name: dict(registry) for name, registry in _codec_registries.items()} + + _register_builtin_codecs() + + assert before == {name: dict(registry) for name, registry in _codec_registries.items()} + assert get_codec_class("custom_builtin_test") is CustomBytesCodec + + +def test_builtin_codec_bootstrap_preserves_lazy_entrypoints( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from zarr.codecs import BytesCodec + from zarr.registry import Registry, _codec_registries, _register_builtin_codecs + + name = "custom_entrypoint_test" + registry: Registry[BytesCodec] = Registry() + registry.lazy_load_list.append( + EntryPoint(name=name, value="zarr.codecs.bytes:BytesCodec", group="zarr.codecs") + ) + monkeypatch.setitem(_codec_registries, name, registry) + + _register_builtin_codecs() + + assert get_codec_class(name) is BytesCodec + assert registry.lazy_load_list == [] + + @pytest.fixture def unregistered_v3_codec(monkeypatch: pytest.MonkeyPatch) -> str: """Guarantee a Zarr format 3 codec name resolves to nothing.