Merge pull request #20 from openclimatefix/jacob/zarr_v3

Add Zarr V3 version based on numcodecs zarr v3

.github/workflows/workflows.yaml

@@ -13,4 +13,4 @@ jobs:
       # pytest-cov looks at this folder
       pytest_cov_dir: "ocf_blosc2"
       os_list: '["ubuntu-latest"]'
-      python-version: "['3.10','3.11']"
+      python-version: "['3.11','3.12']"

.ruff.toml

@@ -1,13 +1,13 @@
 # Enable pycodestyle (`E`) and Pyflakes (`F`) codes by default.
-select = ["E", "F", "D", "I"]
-ignore = ["D200","D202","D210","D212","D415","D105",]
+lint.select = ["E", "F", "D", "I"]
+lint.ignore = ["D200","D202","D210","D212","D415","D105",]
 
 # Allow autofix for all enabled rules (when `--fix`) is provided.
-fixable = ["A", "B", "C", "D", "E", "F", "I"]
-unfixable = []
+lint.fixable = ["A", "B", "C", "D", "E", "F", "I"]
+lint.unfixable = []
 
 # Exclude a variety of commonly ignored directories.
-exclude = [
+lint.exclude = [
     ".bzr",
     ".direnv",
     ".eggs",
@@ -35,21 +35,19 @@ exclude = [
 line-length = 100
 
 # Allow unused variables when underscore-prefixed.
-dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+lint.dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 
 # Assume Python 3.10.
 target-version = "py310"
 fix = false
-# Group violations by containing file.
-ignore-init-module-imports = true
 
-[mccabe]
+[lint.mccabe]
 # Unlike Flake8, default to a complexity level of 10.
 max-complexity = 10
 
-[pydocstyle]
+[lint.pydocstyle]
 # Use Google-style docstrings.
 convention = "google"
 
-[per-file-ignores]
+[lint.per-file-ignores]
 "__init__.py" = ["F401", "E402"]

ocf_blosc2/ocf_blosc2_v3.py

@@ -0,0 +1,5 @@
+"""Zarr V3 adaption of Blosc2 codec for numcodecs."""
+
+from numcodecs.zarr3 import _make_bytes_bytes_codec
+
+Blosc2 = _make_bytes_bytes_codec("blosc2", "Blosc2")

requirements.txt

@@ -1,4 +1,5 @@
-numpy==1.26.4
+numpy
 numcodecs
 blosc2
 pytest
+zarr>=3.0.0

setup.py

@@ -17,7 +17,7 @@
     author="Jacob Bieker",
     author_email="[email protected]",
     company="Open Climate Fix Ltd",
-    install_requires=["numpy", "blosc2", "numcodecs"],
+    install_requires=["numpy", "blosc2", "numcodecs", "zarr>3"],
     long_description=long_description,
     long_description_content_type="text/markdown",
     packages=find_packages(),

tests/test_ocf_blosc2.py

@@ -6,6 +6,7 @@
 
 
 def test_roundtrip():
+    """Test roundtrip encoding and decoding of a numpy array."""
     buf = np.asarray([np.nan, 0.0, 0.5, 1.0], dtype=np.float32)
     blosc2 = ocf_blosc2.Blosc2()
     comp_arr = blosc2.encode(buf)

tests/test_zarr3.py

@@ -0,0 +1,69 @@
+"""Test the Zarr version 3 codec with the Blosc2 codec."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pytest
+
+if TYPE_CHECKING:  # pragma: no cover
+    import zarr
+else:
+    zarr = pytest.importorskip("zarr")
+
+import numcodecs.zarr3
+import zarr.storage
+
+from ocf_blosc2.ocf_blosc2_v3 import Blosc2
+
+pytestmark = [
+    pytest.mark.filterwarnings("ignore:Codec 'numcodecs.*' not configured in config.*:UserWarning"),
+    pytest.mark.filterwarnings(
+        "ignore:Numcodecs codecs are not in the Zarr version 3 specification and may not be supported by other zarr implementations."  # noqa
+    ),
+]
+
+get_codec_class = zarr.registry.get_codec_class
+Array = zarr.Array
+BytesCodec = zarr.codecs.BytesCodec
+Store = zarr.abc.store.Store
+MemoryStore = zarr.storage.MemoryStore
+StorePath = zarr.storage.StorePath
+
+
+EXPECTED_WARNING_STR = "Numcodecs codecs are not in the Zarr version 3.*"
+
+
+@pytest.fixture
+def store() -> StorePath:
+    """Make a new in-memory store."""
+    return StorePath(MemoryStore(read_only=False))
+
+
+ALL_CODECS = [getattr(numcodecs.zarr3, cls_name) for cls_name in numcodecs.zarr3.__all__]
+
+
+@pytest.mark.parametrize(
+    "codec_class",
+    [Blosc2],
+)
+def test_generic_compressor(
+    store: StorePath, codec_class: type[numcodecs.zarr3._NumcodecsBytesBytesCodec]
+):
+    """Test that the generic compressor works with the Blosc2 codec."""
+    data = np.arange(0, 256, dtype="uint16").reshape((16, 16))
+
+    with pytest.warns(UserWarning, match=EXPECTED_WARNING_STR):
+        a = zarr.create_array(
+            store / "generic",
+            shape=data.shape,
+            chunks=(4, 4),
+            shards=(16, 16),
+            dtype=data.dtype,
+            fill_value=0,
+            compressors=[codec_class()],
+        )
+
+    a[:, :] = data.copy()
+    np.testing.assert_array_equal(data, a[:, :])