Add support for zarr v3

.github/workflows/python-package-binaries.yml

@@ -14,7 +14,7 @@ jobs:
       fail-fast: false
       matrix:
         os: ["ubuntu-22.04", "ubuntu-24.04"]
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.10", "3.13"]
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5

.github/workflows/python-package-cython.yml

@@ -8,13 +8,14 @@ on:
 
 jobs:
   build-and-test:
-    name: Test on ${{ matrix.os }} - Python ${{ matrix.python-version }}
+    name: Test on ${{ matrix.os }} - Python ${{ matrix.python-version }} - Zarr ${{ matrix.zarr-version }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         os: ["ubuntu-latest", "macos-latest"]
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.11", "3.13"]
+        zarr-version: ["2.18.4", "3.1.5"]
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
@@ -45,6 +46,9 @@ jobs:
       - name: Install dependencies
         run: |
           pip install .[test]
+      - name: Install Zarr
+        run: |
+          pip install zarr==${{ matrix.zarr-version }}
       - name: Test imports and version
         run: |
           pytest -s tests/test_imports.py

.github/workflows/python-package-multi-threading.yml

@@ -14,7 +14,7 @@ jobs:
       fail-fast: false
       matrix:
         os: ["ubuntu-latest", "macos-latest"]
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.10", "3.13"]
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5

.github/workflows/python-package-no-cython.yml

@@ -14,7 +14,7 @@ jobs:
       fail-fast: false
       matrix:
         os: ["ubuntu-latest", "macos-latest", "windows-latest"]
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.10", "3.13"]
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5

README.md

@@ -68,16 +68,24 @@ data = ... # any numpy array
 # instantiate WavPack compressor
 wv_compressor = WavPack(level=2, bps=None)
 
-z = zarr.array(data, compressor=wv_compressor)
+# v2
+group = zarr.group()
+z = group.create_dataset(name="wv_dset", data=data, compressor=wv_compressor)
+
+# v3
+zarr.config.set({"default_zarr_version": 3})
+group = zarr.group()
+z = group.create(name="wv_dset3", data=data, codecs=[wv_compressor])
 
 data_read = z[:]
 ```
 Available `**kwargs` can be browsed with: `WavPack?`
 
-**NOTE:** 
-In order to reload in zarr an array saved with the `WavPack`, you just need to have the `wavpack_numcodecs` package
+> **_NOTE 1:_**  In order to reload in zarr an array saved with the `WavPack`, you just need to have the `wavpack_numcodecs` package
 installed.
 
+> **_NOTE 2:_** The Zarr v3 implementation is an `ArrayBytesCodec`. The `zarr.create_array` function only supports `ArrayArrayCodec` objects for `filters` and `BytesBytesCodecs` for `compressors`. Hence, we need to use the `zarr.create` function instead, which support any list of codecs (including `ArrayBytesCodec` objects).
+
 # Developmers guide
 
 ## How to upgrade WavPack installation and make a new release

src/wavpack_numcodecs/__init__.py

@@ -1,4 +1,28 @@
-from wavpack_numcodecs.wavpack import WavPack, wavpack_version
+import importlib.metadata
+import importlib.util
+import numcodecs
+from packaging.version import parse
+
+from wavpack_numcodecs.wavpack import wavpack_version
+
+
+HAVE_ZARR = importlib.util.find_spec("zarr") is not None
+
+USE_ZARR_V3 = False
+if HAVE_ZARR:
+    import zarr
+
+    if parse(zarr.__version__) >= parse("3.0.0"):
+        USE_ZARR_V3 = True
+
+if USE_ZARR_V3:
+    from zarr.registry import register_codec
+    from wavpack_numcodecs.wavpackv3 import WavPack
+else:
+    from numcodecs import register_codec
+    from wavpack_numcodecs.wavpack import WavPack
+
+register_codec("wavpack", WavPack)
 
 from .globals import (
     get_num_decoding_threads,
@@ -8,6 +32,5 @@
     set_num_decoding_threads,
     set_num_encoding_threads,
 )
-import importlib.metadata
 
-__version__ = importlib.metadata.version("wavpack_numcodecs")
+__version__ = importlib.metadata.version("wavpack_numcodecs")

src/wavpack_numcodecs/wavpack.pyx

@@ -329,6 +329,3 @@ class WavPack(Codec):
     def decode(self, buf, out=None):        
         buf = ensure_contiguous_ndarray(buf, self.max_buffer_size)
         return decompress(buf, out, self.num_decoding_threads)
-
-
-numcodecs.register_codec(WavPack)

src/wavpack_numcodecs/wavpackv3.py

@@ -0,0 +1,97 @@
+from zarr.abc.codec import ArrayBytesCodec
+from zarr.core.buffer import Buffer, BufferPrototype
+from zarr.core.common import BytesLike
+from wavpack_numcodecs.wavpack import WavPack as WavPackV2
+import numpy as np
+import asyncio
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from zarr.core.array_spec import ArraySpec
+
+
+class WavPack(ArrayBytesCodec):
+    def __init__(
+        self,
+        level: int = 1,
+        bps: int | None = None,
+        dynamic_noise_shaping: bool = True,
+        shaping_weight: float = 0.0,
+        num_encoding_threads: int = 1,
+        num_decoding_threads: int = 8,
+    ):
+        self._codec = WavPackV2(
+            level=level,
+            bps=bps,
+            dynamic_noise_shaping=dynamic_noise_shaping,
+            shaping_weight=shaping_weight,
+            num_encoding_threads=num_encoding_threads,
+            num_decoding_threads=num_decoding_threads,
+        )
+
+    async def _encode_single(
+        self,
+        chunk_array: np.ndarray,
+        chunk_spec: "ArraySpec",
+    ) -> Buffer | None:
+        """Encode a single chunk."""
+        # Convert to numpy array if it's an NDBuffer
+        if hasattr(chunk_array, "as_numpy_array"):
+            chunk_array = chunk_array.as_numpy_array()
+        elif not isinstance(chunk_array, np.ndarray):
+            chunk_array = np.asarray(chunk_array)
+
+        encoded = await asyncio.to_thread(self._codec.encode, chunk_array)
+        return chunk_spec.prototype.buffer.from_bytes(encoded)
+
+    async def _decode_single(
+        self,
+        chunk_bytes: Buffer,
+        chunk_spec: "ArraySpec",
+    ) -> np.ndarray:
+        """Decode a single chunk."""
+        decoded = await asyncio.to_thread(self._codec.decode, chunk_bytes.to_bytes())
+
+        # Convert to numpy array if it's bytes
+        if isinstance(decoded, bytes):
+            np_dtype = chunk_spec.dtype.to_native_dtype()
+            decoded = np.frombuffer(decoded, dtype=np_dtype)
+
+        # Ensure it's a numpy array with correct shape
+        if isinstance(decoded, np.ndarray):
+            return decoded.reshape(chunk_spec.shape)
+        else:
+            raise TypeError(f"Expected numpy array from decode, got {type(decoded)}")
+
+    def compute_encoded_size(self, input_byte_length: int, chunk_spec: "ArraySpec") -> int:
+        # WavPack compression ratio is variable, so we can't predict the exact size
+        # Return a conservative estimate
+        return input_byte_length
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "WavPack":
+        """Create codec from configuration dictionary."""
+        config = data.get("configuration", {})
+        return cls(
+            level=config.get("level", 1),
+            bps=config.get("bps"),
+            dynamic_noise_shaping=config.get("dynamic_noise_shaping", True),
+            shaping_weight=config.get("shaping_weight", 0.0),
+            num_encoding_threads=config.get("num_encoding_threads", 1),
+            num_decoding_threads=config.get("num_decoding_threads", 8),
+        )
+
+    def to_dict(self) -> dict:
+        """Convert codec to configuration dictionary."""
+        config = self._codec.get_config()
+        return {
+            "name": "wavpack",
+            "configuration": {
+                "level": config.get("level", 1),
+                "bps": config.get("bps"),
+                "dynamic_noise_shaping": config.get("dynamic_noise_shaping", True),
+                "shaping_weight": config.get("shaping_weight", 0.0),
+                "num_encoding_threads": config.get("num_encoding_threads", 1),
+                "num_decoding_threads": config.get("num_decoding_threads", 8),
+            },
+        }

tests/test_imports.py

@@ -2,7 +2,7 @@ def test_imports():
     from wavpack_numcodecs import wavpack_version
 
     print(f"\nWavpack library verison: {wavpack_version}")
-    from wavpack_numcodecs import WavPack
+    from wavpack_numcodecs.wavpack import WavPack
 
     wv0 = WavPack(level=2)
     print(wv0)
@@ -11,8 +11,8 @@ def test_imports():
 
 
 def test_global_settings():
+    from wavpack_numcodecs.wavpack import WavPack
     from wavpack_numcodecs import (
-        WavPack,
         get_num_decoding_threads,
         get_num_encoding_threads,
         reset_num_decoding_threads,