Backward and forward-compatible changes to support zarr-python v3

Author: tomwhite

dev.py

@@ -56,7 +56,7 @@ def main(root):
     # we avoid retrieving stuff we don't need.
     format_fields = {}
     info_fields = {}
-    for name, array in root.items():
+    for name, array in root.arrays():
         if name.startswith("call_") and not name.startswith("call_genotype"):
             vcf_name = name[len("call_") :]
             format_fields[vcf_name] = array.blocks[v_chunk]

vcztools/filter.py

@@ -113,7 +113,7 @@ def _compose_comparison_evaluator(self, parse_results: pp.ParseResults) -> Calla
 
         def evaluator(root, variant_chunk_index: int) -> np.ndarray:
             vcf_name = parse_results[0]
-            vcz_names = set(name for name, _array in root.items())
+            vcz_names = set(root.keys())
             vcz_name = vcf_name_to_vcz_name(vcz_names, vcf_name)
             zarray = root[vcz_name]
             variant_chunk_len = zarray.chunks[0]

vcztools/query.py

@@ -130,7 +130,7 @@ def _compose_tag_generator(
             return self._compose_gt_generator()
 
         def generate(root):
-            vcz_names = set(name for name, _zarray in root.items())
+            vcz_names = set(root.keys())
             vcz_name = vcf_name_to_vcz_name(vcz_names, tag)
             zarray = root[vcz_name]
             contig_ids = root["contig_id"][:] if tag == "CHROM" else None

vcztools/regions.py

@@ -47,10 +47,14 @@ def create_index(vcz) -> None:
     root.array(
         "region_index",
         data=index,
+        shape=index.shape,
+        dtype=index.dtype,
         compressor=numcodecs.Blosc("zstd", clevel=9, shuffle=0),
-        overwrite=True,
     )
 
+    # consolidate to ensure region_index is found
+    zarr.consolidate_metadata(vcz)
+
 
 def parse_region_string(region: str) -> tuple[str, int | None, int | None]:
     """Return the contig, start position and end position from a region string."""

vcztools/vcf_writer.py

@@ -292,17 +292,17 @@ def c_chunk_to_vcf(
     drop_genotypes,
     no_update,
 ):
-    chrom = contigs[get_vchunk_array(root.variant_contig, v_chunk, v_mask_chunk)]
+    chrom = contigs[get_vchunk_array(root["variant_contig"], v_chunk, v_mask_chunk)]
     # TODO check we don't truncate silently by doing this
-    pos = get_vchunk_array(root.variant_position, v_chunk, v_mask_chunk).astype(
+    pos = get_vchunk_array(root["variant_position"], v_chunk, v_mask_chunk).astype(
         np.int32
     )
-    id = get_vchunk_array(root.variant_id, v_chunk, v_mask_chunk).astype("S")
-    alleles = get_vchunk_array(root.variant_allele, v_chunk, v_mask_chunk)
+    id = get_vchunk_array(root["variant_id"], v_chunk, v_mask_chunk).astype("S")
+    alleles = get_vchunk_array(root["variant_allele"], v_chunk, v_mask_chunk)
     ref = alleles[:, 0].astype("S")
     alt = alleles[:, 1:].astype("S")
-    qual = get_vchunk_array(root.variant_quality, v_chunk, v_mask_chunk)
-    filter_ = get_vchunk_array(root.variant_filter, v_chunk, v_mask_chunk)
+    qual = get_vchunk_array(root["variant_quality"], v_chunk, v_mask_chunk)
+    filter_ = get_vchunk_array(root["variant_filter"], v_chunk, v_mask_chunk)
 
     num_variants = len(pos)
     if len(id.shape) == 1:
@@ -313,7 +313,7 @@ def c_chunk_to_vcf(
     format_fields = {}
     info_fields = {}
     num_samples = len(samples_selection) if samples_selection is not None else None
-    for name, array in root.items():
+    for name, array in root.arrays():
         if (
             name.startswith("call_")
             and not name.startswith("call_genotype")
@@ -417,7 +417,8 @@ def _generate_header(ds, original_header, sample_ids, *, no_version: bool = Fals
         # GT must be the first field if present, per the spec (section 1.6.2)
         format_fields.append("GT")
 
-    for var, arr in ds.items():
+    for var in sorted(ds.keys()):
+        arr = ds[var]
         if (
             var.startswith("variant_")
             and not var.endswith("_fill")