Switch Zarr DirectoryStore to zarr.open() for v2/v3 compat

Author: jeromekelleher

bio2zarr/plink.py

@@ -17,8 +17,7 @@ def encode_genotypes_slice(bed_path, zarr_path, start, stop):
     # the correct approach is, but it is important to note that the
     # 0th allele is *not* necessarily the REF for these datasets.
     bed = bed_reader.open_bed(bed_path, num_threads=1, count_A1=False)
-    store = zarr.DirectoryStore(zarr_path)
-    root = zarr.group(store=store)
+    root = zarr.open(store=zarr_path, mode="a")
     gt = core.BufferedArray(root["call_genotype"], start)
     gt_mask = core.BufferedArray(root["call_genotype_mask"], start)
     gt_phased = core.BufferedArray(root["call_genotype_phased"], start)
@@ -73,8 +72,7 @@ def convert(
     if variants_chunk_size is None:
         variants_chunk_size = 10_000
 
-    store = zarr.DirectoryStore(zarr_path)
-    root = zarr.group(store=store, overwrite=True)
+    root = zarr.open(store=zarr_path, mode="w")
 
     ploidy = 2
     shape = [m, n]
@@ -171,8 +169,7 @@ def convert(
 # FIXME do this more efficiently - currently reading the whole thing
 # in for convenience, and also comparing call-by-call
 def validate(bed_path, zarr_path):
-    store = zarr.DirectoryStore(zarr_path)
-    root = zarr.group(store=store)
+    root = zarr.open(store=zarr_path, mode="r")
     call_genotype = root["call_genotype"][:]
 
     bed = bed_reader.open_bed(bed_path, count_A1=False, num_threads=1)

bio2zarr/vcf2zarr/vcz.py

@@ -532,8 +532,7 @@ def init(
         )
 
         self.path.mkdir()
-        store = zarr.DirectoryStore(self.path)
-        root = zarr.group(store=store)
+        root = zarr.open(store=self.path, mode="a")
         root.attrs.update(
             {
                 "vcf_zarr_version": "0.2",
@@ -549,8 +548,7 @@ def init(
         self.wip_path.mkdir()
         self.arrays_path.mkdir()
         self.partitions_path.mkdir()
-        store = zarr.DirectoryStore(self.arrays_path)
-        root = zarr.group(store=store)
+        root = zarr.open(store=self.arrays_path, mode="a")
 
         total_chunks = 0
         for field in self.schema.fields:
@@ -690,8 +688,7 @@ def init_partition_array(self, partition_index, name):
         # Overwrite any existing WIP files
         wip_path = self.wip_partition_array_path(partition_index, name)
         shutil.copytree(src, wip_path, dirs_exist_ok=True)
-        store = zarr.DirectoryStore(self.wip_partition_path(partition_index))
-        wip_root = zarr.group(store=store)
+        wip_root = zarr.open(store=self.wip_partition_path(partition_index), mode="a")
         array = wip_root[name]
         logger.debug(f"Opened empty array {array.name} <{array.dtype}> @ {wip_path}")
         return array

bio2zarr/vcf2zarr/verification.py

@@ -145,9 +145,7 @@ def assert_format_val_equal(vcf_val, zarr_val, vcf_type, vcf_number):
 
 
 def verify(vcf_path, zarr_path, show_progress=False):
-    store = zarr.DirectoryStore(zarr_path)
-
-    root = zarr.group(store=store)
+    root = zarr.open(store=zarr_path, mode="r")
     pos = root["variant_position"][:]
     allele = root["variant_allele"][:]
     chrom = root["contig_id"][:][root["variant_contig"][:]]