refactor: centralize built-in codec registration

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"):

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")

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)
 

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.