Add zarr v3 support

Author: benjeffery

.github/workflows/tests.yml

@@ -50,5 +50,52 @@ jobs:
       - name: Upload coverage to Codecov
         uses: codecov/[email protected]
         with:
+          name: zarrv3
           token: ${{ secrets.CODECOV_TOKEN }}
           fail_ci_if_error: true
+
+      - name: Run tests (zarr v2)
+        if: ${{ matrix.python == '3.13' }}
+        run: |
+          uv pip install "zarr>=2.17,<3"
+          uv run --no-sync python -m pytest -xv --cov=tszip --cov-report=xml --cov-branch -n2
+
+      - name: Upload coverage to Codecov
+        uses: codecov/[email protected]
+        with:
+          name: zarrv2
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: true
+
+  zarr-compatibility:
+    name: Zarr v2/v3 Cross-compatibility
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/[email protected]
+
+      - name: Install uv and set python version
+        uses: astral-sh/setup-uv@v6
+        with:
+          python-version: 3.13
+          version: "0.8.15"
+
+      - name: Install dependencies (zarr v3)
+        run: |
+          uv venv
+          uv pip install -r pyproject.toml --extra test
+
+      - name: Write test file with zarr v3
+        run: uv run --no-sync python tests/zarr_cross_version_helper.py write test_v3.tsz
+
+      - name: Switch to zarr v2 and test reading
+        run: |
+          uv pip install "zarr>=2.17,<3"
+          uv run --no-sync python tests/zarr_cross_version_helper.py read test_v3.tsz
+          uv run --no-sync python tests/zarr_cross_version_helper.py write test_v2.tsz
+
+      - name: Switch back to zarr v3 and test reading both files
+        run: |
+          uv pip install "zarr>=3.0,<4"
+          uv run --no-sync python tests/zarr_cross_version_helper.py read test_v3.tsz
+          uv run --no-sync python tests/zarr_cross_version_helper.py read test_v2.tsz

CHANGELOG.rst

@@ -3,6 +3,7 @@
 --------------------
 
 - Drop Python 3.9 support, require Python >= 3.10 (#112, benjeffery)
+- Support zarr v3 (#114, benjeffery)
 
 --------------------
 [0.2.4] - 2024-07-10

pyproject.toml

@@ -35,7 +35,7 @@ dependencies = [
     "humanize",
     "tskit>=0.3.3",
     "numcodecs>=0.6.4",
-    "zarr<3",
+    "zarr>=2.17,<4",
 ]
 dynamic = ["version"]
 
@@ -59,7 +59,8 @@ test = [
     "pytest-cov==6.3.0",
     "pytest-xdist==3.8.0",
     "tskit==0.6.4",
-    "zarr==2.17.2",
+    "zarr==2.18.3; python_version == '3.10'",
+    "zarr==3.1.2; python_version >= '3.11'",
     "numcodecs>=0.6,<0.15.1", #Pinned due to https://github.com/zarr-developers/numcodecs/issues/733
 ]
 docs = [
@@ -69,7 +70,7 @@ docs = [
     "sphinx-argparse==0.5.2",
     "setuptools_scm==9.2.0",
     "tskit==0.6.4",
-    "zarr==2.18.7",
+    "zarr==3.1.2",
     "numcodecs>=0.6,<0.15.1", #Pinned due to https://github.com/zarr-developers/numcodecs/issues/733
 ]
 dev = [
@@ -87,7 +88,7 @@ dev = [
     "sphinx-issues",
     "setuptools_scm",
     "tskit",
-    "zarr<3",
+    "zarr",
     "msprime",
     "humanize",
 ]

tests/test_compression.py

@@ -30,12 +30,12 @@
 import numpy as np
 import pytest
 import tskit
-import zarr
 
 import tszip
 import tszip.compression as compression
 import tszip.exceptions as exceptions
 import tszip.provenance as provenance
+from tszip import compat
 
 
 class TestMinimalDtype(unittest.TestCase):
@@ -294,25 +294,25 @@ def tearDown(self):
     def test_format_written(self):
         ts = msprime.simulate(10, random_seed=1)
         tszip.compress(ts, self.path)
-        with zarr.ZipStore(str(self.path), mode="r") as store:
-            root = zarr.group(store=store)
+        with compat.create_zip_store(str(self.path), mode="r") as store:
+            root = compat.create_zarr_group(store=store)
             self.assertEqual(root.attrs["format_name"], compression.FORMAT_NAME)
             self.assertEqual(root.attrs["format_version"], compression.FORMAT_VERSION)
 
     def test_provenance(self):
         ts = msprime.simulate(10, random_seed=1)
         for variants_only in [True, False]:
             tszip.compress(ts, self.path, variants_only=variants_only)
-            with zarr.ZipStore(str(self.path), mode="r") as store:
-                root = zarr.group(store=store)
+            with compat.create_zip_store(str(self.path), mode="r") as store:
+                root = compat.create_zarr_group(store=store)
                 self.assertEqual(
                     root.attrs["provenance"],
                     provenance.get_provenance_dict({"variants_only": variants_only}),
                 )
 
     def write_file(self, attrs, path):
-        with zarr.ZipStore(str(path), mode="w") as store:
-            root = zarr.group(store=store)
+        with compat.create_zip_store(str(path), mode="w") as store:
+            root = compat.create_zarr_group(store=store)
             root.attrs.update(attrs)
 
     def test_missing_format_keys(self):

tests/zarr_cross_version_helper.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+"""
+Script to test zarr cross-version compatibility.
+Usage: python test_zarr_cross_version.py [write|read] <filename>
+"""
+import pathlib
+import sys
+
+import msprime
+import tskit
+
+# Add parent directory to path so we can import tszip
+sys.path.insert(0, str(pathlib.Path(__file__).parent.parent))
+
+import tszip  # noqa: E402
+
+
+def all_fields_ts(edge_metadata=True, migrations=True):
+    """
+    A tree sequence with data in all fields (except edge metadata is not set if
+    edge_metadata is False and migrations are not defined if migrations is False
+    (this is needed to test simplify, which doesn't allow either)
+
+    """
+    demography = msprime.Demography()
+    demography.add_population(name="A", initial_size=10_000)
+    demography.add_population(name="B", initial_size=5_000)
+    demography.add_population(name="C", initial_size=1_000)
+    demography.add_population(name="D", initial_size=500)
+    demography.add_population(name="E", initial_size=100)
+    demography.add_population_split(time=1000, derived=["A", "B"], ancestral="C")
+    ts = msprime.sim_ancestry(
+        samples={"A": 10, "B": 10},
+        demography=demography,
+        sequence_length=5,
+        random_seed=42,
+        recombination_rate=1,
+        record_migrations=migrations,
+        record_provenance=True,
+    )
+    ts = msprime.sim_mutations(ts, rate=0.001, random_seed=42)
+    tables = ts.dump_tables()
+    # Add locations to individuals
+    individuals_copy = tables.individuals.copy()
+    tables.individuals.clear()
+    for i, individual in enumerate(individuals_copy):
+        tables.individuals.append(
+            individual.replace(flags=i, location=[i, i + 1], parents=[i - 1, i - 1])
+        )
+    # Ensure all columns have unique values
+    nodes_copy = tables.nodes.copy()
+    tables.nodes.clear()
+    for i, node in enumerate(nodes_copy):
+        tables.nodes.append(
+            node.replace(
+                flags=i,
+                time=node.time + 0.00001 * i,
+                individual=i % len(tables.individuals),
+                population=i % len(tables.populations),
+            )
+        )
+    if migrations:
+        tables.migrations.add_row(left=0, right=1, node=21, source=1, dest=3, time=1001)
+
+    # Add metadata
+    for name, table in tables.table_name_map.items():
+        if name == "provenances":
+            continue
+        if name == "migrations" and not migrations:
+            continue
+        if name == "edges" and not edge_metadata:
+            continue
+        table.metadata_schema = tskit.MetadataSchema.permissive_json()
+        metadatas = [f'{{"foo":"n_{name}_{u}"}}' for u in range(len(table))]
+        metadata, metadata_offset = tskit.pack_strings(metadatas)
+        table.set_columns(
+            **{
+                **table.asdict(),
+                "metadata": metadata,
+                "metadata_offset": metadata_offset,
+            }
+        )
+    tables.metadata_schema = tskit.MetadataSchema.permissive_json()
+    tables.metadata = "Test metadata"
+    tables.time_units = "Test time units"
+
+    tables.reference_sequence.metadata_schema = tskit.MetadataSchema.permissive_json()
+    tables.reference_sequence.metadata = "Test reference metadata"
+    tables.reference_sequence.data = "A" * int(ts.sequence_length)
+    tables.reference_sequence.url = "http://example.com/a_reference"
+
+    # Add some more rows to provenance to have enough for testing.
+    for i in range(3):
+        tables.provenances.add_row(record="A", timestamp=str(i))
+
+    return tables.tree_sequence()
+
+
+def write_test_file(filename):
+    """Write a test file with current zarr version"""
+    ts = all_fields_ts()
+    tszip.compress(ts, filename)
+    ts2 = tszip.decompress(filename)
+    ts.tables.assert_equals(ts2.tables)
+
+
+def read_test_file(filename):
+    """Read and verify a test file with current zarr version"""
+    try:
+        tszip.decompress(filename)
+    except Exception:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    action = sys.argv[1]
+    filename = sys.argv[2]
+    if action == "write":
+        write_test_file(filename)
+    elif action == "read":
+        read_test_file(filename)

tszip/compat.py

@@ -0,0 +1,93 @@
+# MIT License
+#
+# Copyright (c) 2025 Tskit Developers
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""
+Compatibility layer for zarr v2/v3 API differences
+"""
+import zarr
+
+ZARR_V3 = zarr.__version__.startswith("3.")
+
+
+if ZARR_V3:
+    from zarr.storage import ZipStore
+
+    def create_zip_store(path, mode="r"):
+        return ZipStore(path, mode=mode)
+
+    def create_zarr_group(store=None):
+        if store is None:
+            return zarr.create_group(zarr_format=2)
+        else:
+            mode = "r" if getattr(store, "read_only", False) else "a"
+            return zarr.open_group(store=store, zarr_format=2, mode=mode)
+
+    def create_empty_array(
+        group, name, shape, dtype, chunks=None, filters=None, compressor=None
+    ):
+        return group.empty(
+            name=name,
+            shape=shape,
+            dtype=dtype,
+            chunks=chunks,
+            zarr_format=2,
+            filters=filters,
+            compressor=compressor,
+        )
+
+    def get_nbytes_stored(array):
+        return array.nbytes_stored()
+
+    def group_items(group):
+        return group.members()
+
+    def visit_arrays(group, visitor):
+        for array in group.array_values():
+            visitor(array)
+
+else:
+
+    def create_zip_store(path, mode="r"):
+        return zarr.ZipStore(path, mode=mode)
+
+    def create_zarr_group(store=None):
+        return zarr.group(store=store)
+
+    def create_empty_array(
+        group, name, shape, dtype, chunks=None, filters=None, compressor=None
+    ):
+        return group.empty(
+            name,
+            shape=shape,
+            dtype=dtype,
+            chunks=chunks,
+            filters=filters,
+            compressor=compressor,
+        )
+
+    def get_nbytes_stored(array):
+        return array.nbytes_stored
+
+    def group_items(group):
+        return group.items()
+
+    def visit_arrays(group, visitor):
+        group.visitvalues(visitor)