format, docstring, and typing

Author: brokkoli71

src/zarr/registry.py

@@ -7,7 +7,6 @@
 
 import numcodecs
 
-from zarr.abc.codec import Codec
 from zarr.core.config import BadConfigError, config
 
 if TYPE_CHECKING:
@@ -17,6 +16,7 @@
         ArrayArrayCodec,
         ArrayBytesCodec,
         BytesBytesCodec,
+        Codec,
         CodecPipeline,
     )
     from zarr.core.buffer import Buffer, NDBuffer
@@ -168,17 +168,18 @@ def _resolve_codec(data: dict[str, JSON]) -> Codec:
     return get_codec_class(data["name"]).from_dict(data)  # type: ignore[arg-type]
 
 
-def numcodec_to_zarr3_codec(codec: numcodecs.abc.Codec) -> Codec:
+def numcodec_to_zarr3_codec(codec: numcodecs.abc.Codec) -> numcodecs.zarr3._NumcodecsCodec & Codec:
+    """
+    Convert a numcodecs codec to a zarr v3 compatible numcodecs.zarr3 codec instance.
+    """
     codec_config = codec.get_config()
     codec_name = codec_config.pop("id", None)
     if codec_name is None:
         raise ValueError(f"Codec configuration does not contain 'id': {codec_config}")
     codec_cls = get_codec_class(f"numcodecs.{codec_name}")
     if codec_cls is None:
         raise ValueError(f"Codec class for 'numcodecs.{codec_name}' not found.")
-    codec_v3 = codec_cls.from_dict({"name": f"numcodecs.{codec_name}", "configuration": codec_config})
-    assert isinstance(codec_v3, Codec)
-    return codec_v3
+    return codec_cls.from_dict({"name": f"numcodecs.{codec_name}", "configuration": codec_config})
 
 
 def _parse_bytes_bytes_codec(data: dict[str, JSON] | Codec) -> BytesBytesCodec:

tests/test_codecs/test_codecs.py

@@ -26,6 +26,7 @@
 from zarr.storage import StorePath
 
 if TYPE_CHECKING:
+    from zarr.abc.codec import Codec
     from zarr.abc.store import Store
     from zarr.core.array import CompressorsLike, FiltersLike, SerializerLike
     from zarr.core.buffer.core import NDArrayLikeOrScalar
@@ -452,7 +453,11 @@ async def test_resize(store: Store) -> None:
         (numcodecs.Zstd(), numcodecs.zarr3.Zstd),
     ],
 )
-def test_numcodecs_in_v3(store: Store, codec_v2, expected_v3_cls) -> None:
+def test_numcodecs_in_v3(
+    store: Store, codec_v2: numcodecs.abc.Codec, expected_v3_cls: Codec
+) -> None:
+    import zarr.registry
+
     result_v3 = zarr.registry.numcodec_to_zarr3_codec(codec_v2)
 
     assert result_v3.__class__ == expected_v3_cls