Clean up zarr.codecs namespace

src/zarr/codecs/__init__.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
-from zarr.codecs.blosc import BloscCname, BloscCodec, BloscShuffle
-from zarr.codecs.bytes import BytesCodec, Endian
-from zarr.codecs.crc32c_ import Crc32cCodec
-from zarr.codecs.gzip import GzipCodec
-from zarr.codecs.pipeline import BatchedCodecPipeline
-from zarr.codecs.sharding import ShardingCodec, ShardingCodecIndexLocation
-from zarr.codecs.transpose import TransposeCodec
-from zarr.codecs.zstd import ZstdCodec
+from zarr.codecs._blosc import BloscCname, BloscCodec, BloscShuffle
+from zarr.codecs._bytes import BytesCodec, Endian
+from zarr.codecs._crc32c_ import Crc32cCodec
+from zarr.codecs._gzip import GzipCodec
+from zarr.codecs._pipeline import BatchedCodecPipeline
+from zarr.codecs._sharding import ShardingCodec, ShardingCodecIndexLocation
+from zarr.codecs._transpose import TransposeCodec
+from zarr.codecs._zstd import ZstdCodec
+from zarr.registry import register_codec, register_pipeline
 
 __all__ = [
     "BatchedCodecPipeline",
@@ -23,3 +24,15 @@
     "TransposeCodec",
     "ZstdCodec",
 ]
+
+
+register_codec("blosc", BloscCodec)
+register_codec("bytes", BytesCodec)
+# compatibility with earlier versions of ZEP1
+register_codec("endian", BytesCodec)
+register_codec("crc32c", Crc32cCodec)
+register_codec("gzip", GzipCodec)
+register_pipeline(BatchedCodecPipeline)
+register_codec("sharding_indexed", ShardingCodec)
+register_codec("transpose", TransposeCodec)
+register_codec("zstd", ZstdCodec)

src/zarr/codecs/blosc.py → src/zarr/codecs/_blosc.py

@@ -11,7 +11,6 @@
 from zarr.abc.codec import BytesBytesCodec
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_enum, parse_named_configuration, to_thread
-from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -189,6 +188,3 @@ async def _encode_single(
 
     def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int:
         raise NotImplementedError
-
-
-register_codec("blosc", BloscCodec)

src/zarr/codecs/bytes.py → src/zarr/codecs/_bytes.py

@@ -10,7 +10,6 @@
 from zarr.abc.codec import ArrayBytesCodec
 from zarr.core.buffer import Buffer, NDArrayLike, NDBuffer
 from zarr.core.common import JSON, parse_enum, parse_named_configuration
-from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -119,9 +118,3 @@ async def _encode_single(
 
     def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int:
         return input_byte_length
-
-
-register_codec("bytes", BytesCodec)
-
-# compatibility with earlier versions of ZEP1
-register_codec("endian", BytesCodec)

src/zarr/codecs/crc32c_.py → src/zarr/codecs/_crc32c_.py

@@ -9,7 +9,6 @@
 
 from zarr.abc.codec import BytesBytesCodec
 from zarr.core.common import JSON, parse_named_configuration
-from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -61,6 +60,3 @@ async def _encode_single(
 
     def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int:
         return input_byte_length + 4
-
-
-register_codec("crc32c", Crc32cCodec)

src/zarr/codecs/gzip.py → src/zarr/codecs/_gzip.py

@@ -8,7 +8,6 @@
 from zarr.abc.codec import BytesBytesCodec
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_named_configuration, to_thread
-from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -70,6 +69,3 @@ def compute_encoded_size(
         _chunk_spec: ArraySpec,
     ) -> int:
         raise NotImplementedError
-
-
-register_codec("gzip", GzipCodec)

src/zarr/codecs/pipeline.py → src/zarr/codecs/_pipeline.py

@@ -18,7 +18,7 @@
 from zarr.core.common import JSON, ChunkCoords, concurrent_map, parse_named_configuration
 from zarr.core.config import config
 from zarr.core.indexing import SelectorTuple, is_scalar, is_total_slice
-from zarr.registry import get_codec_class, register_pipeline
+from zarr.registry import get_codec_class
 
 if TYPE_CHECKING:
     import numpy as np
@@ -526,6 +526,3 @@ def codecs_from_list(
         raise ValueError("Required ArrayBytesCodec was not found.")
     else:
         return array_array, array_bytes_maybe, bytes_bytes
-
-
-register_pipeline(BatchedCodecPipeline)

src/zarr/codecs/sharding.py → src/zarr/codecs/_sharding.py

@@ -18,8 +18,8 @@
     CodecPipeline,
 )
 from zarr.abc.store import ByteGetter, ByteSetter
-from zarr.codecs.bytes import BytesCodec
-from zarr.codecs.crc32c_ import Crc32cCodec
+from zarr.codecs._bytes import BytesCodec
+from zarr.codecs._crc32c_ import Crc32cCodec
 from zarr.core.array_spec import ArraySpec
 from zarr.core.buffer import (
     Buffer,
@@ -45,7 +45,7 @@
     morton_order_iter,
 )
 from zarr.core.metadata import parse_codecs
-from zarr.registry import get_ndbuffer_class, get_pipeline_class, register_codec
+from zarr.registry import get_ndbuffer_class, get_pipeline_class
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable, Iterator
@@ -722,6 +722,3 @@ async def _load_full_shard_maybe(
     def compute_encoded_size(self, input_byte_length: int, shard_spec: ArraySpec) -> int:
         chunks_per_shard = self._get_chunks_per_shard(shard_spec)
         return input_byte_length + self._shard_index_size(chunks_per_shard)
-
-
-register_codec("sharding_indexed", ShardingCodec)

src/zarr/codecs/transpose.py → src/zarr/codecs/_transpose.py

@@ -9,7 +9,6 @@
 from zarr.abc.codec import ArrayArrayCodec
 from zarr.core.array_spec import ArraySpec
 from zarr.core.common import JSON, ChunkCoordsLike, parse_named_configuration
-from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing import Any
@@ -109,6 +108,3 @@ async def _encode_single(
 
     def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int:
         return input_byte_length
-
-
-register_codec("transpose", TransposeCodec)

src/zarr/codecs/zstd.py → src/zarr/codecs/_zstd.py

@@ -10,7 +10,6 @@
 from zarr.abc.codec import BytesBytesCodec
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_named_configuration, to_thread
-from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -88,6 +87,3 @@ async def _encode_single(
 
     def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int:
         raise NotImplementedError
-
-
-register_codec("zstd", ZstdCodec)

tests/v3/test_buffer.py

@@ -6,12 +6,7 @@
 import pytest
 
 from zarr import AsyncArray
-from zarr.codecs.blosc import BloscCodec
-from zarr.codecs.bytes import BytesCodec
-from zarr.codecs.crc32c_ import Crc32cCodec
-from zarr.codecs.gzip import GzipCodec
-from zarr.codecs.transpose import TransposeCodec
-from zarr.codecs.zstd import ZstdCodec
+from zarr.codecs import BloscCodec, BytesCodec, Crc32cCodec, GzipCodec, TransposeCodec, ZstdCodec
 from zarr.core.buffer import ArrayLike, BufferPrototype, NDArrayLike, cpu, gpu
 from zarr.store.common import StorePath
 from zarr.store.memory import MemoryStore

tests/v3/test_metadata/test_v3.py

@@ -4,7 +4,7 @@
 import re
 from typing import TYPE_CHECKING, Literal
 
-from zarr.codecs.bytes import BytesCodec
+from zarr.codecs import BytesCodec
 from zarr.core.buffer import default_buffer_prototype
 from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding, V2ChunkKeyEncoding