Address comments

Author: benjeffery

bio2zarr/plink.py

@@ -58,14 +58,19 @@ def generate_schema(
         m = self.bed.sid_count
         logging.info(f"Scanned plink with {n} samples and {m} variants")
 
-        # FIXME
-        if samples_chunk_size is None:
-            samples_chunk_size = 1000
-        if variants_chunk_size is None:
-            variants_chunk_size = 10_000
+        schema_instance = schema.VcfZarrSchema(
+            format_version=schema.ZARR_SCHEMA_FORMAT_VERSION,
+            samples_chunk_size=samples_chunk_size,
+            variants_chunk_size=variants_chunk_size,
+            fields=[],
+            samples=[schema.Sample(id=sample) for sample in self.bed.iid],
+            contigs=[],
+            filters=[],
+        )
 
         logger.info(
-            f"Generating schema with chunks={variants_chunk_size, samples_chunk_size}"
+            "Generating schema with chunks="
+            f"{schema_instance.variants_chunk_size, schema_instance.samples_chunk_size}"
         )
 
         array_specs = [
@@ -75,7 +80,7 @@ def generate_schema(
                 dtype="i4",
                 shape=[m],
                 dimensions=["variants"],
-                chunks=[variants_chunk_size],
+                chunks=[schema_instance.variants_chunk_size],
                 description=None,
             ),
             schema.ZarrArraySpec.new(
@@ -84,7 +89,7 @@ def generate_schema(
                 dtype="O",
                 shape=[m, 2],
                 dimensions=["variants", "alleles"],
-                chunks=[variants_chunk_size, 2],
+                chunks=[schema_instance.variants_chunk_size, 2],
                 description=None,
             ),
             schema.ZarrArraySpec.new(
@@ -93,7 +98,10 @@ def generate_schema(
                 dtype="bool",
                 shape=[m, n],
                 dimensions=["variants", "samples"],
-                chunks=[variants_chunk_size, samples_chunk_size],
+                chunks=[
+                    schema_instance.variants_chunk_size,
+                    schema_instance.samples_chunk_size,
+                ],
                 description=None,
             ),
             schema.ZarrArraySpec.new(
@@ -102,7 +110,11 @@ def generate_schema(
                 dtype="i1",
                 shape=[m, n, 2],
                 dimensions=["variants", "samples", "ploidy"],
-                chunks=[variants_chunk_size, samples_chunk_size, 2],
+                chunks=[
+                    schema_instance.variants_chunk_size,
+                    schema_instance.samples_chunk_size,
+                    2,
+                ],
                 description=None,
             ),
             schema.ZarrArraySpec.new(
@@ -111,20 +123,16 @@ def generate_schema(
                 dtype="bool",
                 shape=[m, n, 2],
                 dimensions=["variants", "samples", "ploidy"],
-                chunks=[variants_chunk_size, samples_chunk_size, 2],
+                chunks=[
+                    schema_instance.variants_chunk_size,
+                    schema_instance.samples_chunk_size,
+                    2,
+                ],
                 description=None,
             ),
         ]
-
-        return schema.VcfZarrSchema(
-            format_version=schema.ZARR_SCHEMA_FORMAT_VERSION,
-            samples_chunk_size=samples_chunk_size,
-            variants_chunk_size=variants_chunk_size,
-            fields=array_specs,
-            samples=[schema.Sample(id=sample) for sample in self.bed.iid],
-            contigs=[],
-            filters=[],
-        )
+        schema_instance.fields = array_specs
+        return schema_instance
 
 
 def convert(

bio2zarr/schema.py

@@ -184,6 +184,28 @@ class VcfZarrSchema(core.JsonDataclass):
     filters: list
     fields: list
 
+    def __init__(
+        self,
+        format_version: str,
+        samples: list,
+        contigs: list,
+        filters: list,
+        fields: list,
+        variants_chunk_size: int = None,
+        samples_chunk_size: int = None,
+    ):
+        self.format_version = format_version
+        self.samples = samples
+        self.contigs = contigs
+        self.filters = filters
+        self.fields = fields
+        if variants_chunk_size is None:
+            variants_chunk_size = 1000
+        self.variants_chunk_size = variants_chunk_size
+        if samples_chunk_size is None:
+            samples_chunk_size = 10_000
+        self.samples_chunk_size = samples_chunk_size
+
     def validate(self):
         """
         Checks that the schema is well-formed and within required limits.

bio2zarr/vcf2zarr/icf.py

@@ -15,9 +15,7 @@
 import numcodecs
 import numpy as np
 
-from bio2zarr import schema
-
-from .. import constants, core, provenance, vcf_utils, writer
+from .. import constants, core, provenance, schema, vcf_utils, writer
 
 logger = logging.getLogger(__name__)
 
@@ -1029,28 +1027,33 @@ def iter_genotypes(self, shape, start, stop):
     def generate_schema(
         self, variants_chunk_size=None, samples_chunk_size=None, local_alleles=None
     ):
-        # Import schema here to avoid circular import
-        from bio2zarr import schema
-
         m = self.num_records
         n = self.num_samples
-        if samples_chunk_size is None:
-            samples_chunk_size = 10_000
-        if variants_chunk_size is None:
-            variants_chunk_size = 1000
         if local_alleles is None:
             local_alleles = False
+
+        schema_instance = schema.VcfZarrSchema(
+            format_version=schema.ZARR_SCHEMA_FORMAT_VERSION,
+            samples_chunk_size=samples_chunk_size,
+            variants_chunk_size=variants_chunk_size,
+            fields=[],
+            samples=self.metadata.samples,
+            contigs=self.metadata.contigs,
+            filters=self.metadata.filters,
+        )
+
         logger.info(
-            f"Generating schema with chunks={variants_chunk_size, samples_chunk_size}"
+            "Generating schema with chunks="
+            f"{schema_instance.variants_chunk_size, schema_instance.samples_chunk_size}"
         )
 
         def spec_from_field(field, array_name=None):
             return schema.ZarrArraySpec.from_field(
                 field,
                 num_samples=n,
                 num_variants=m,
-                samples_chunk_size=samples_chunk_size,
-                variants_chunk_size=variants_chunk_size,
+                samples_chunk_size=schema_instance.samples_chunk_size,
+                variants_chunk_size=schema_instance.variants_chunk_size,
                 array_name=array_name,
             )
 
@@ -1069,7 +1072,7 @@ def fixed_field_spec(
                 shape=shape,
                 description="",
                 dimensions=dimensions,
-                chunks=chunks or [variants_chunk_size],
+                chunks=chunks or [schema_instance.variants_chunk_size],
             )
 
         alt_field = self.fields["ALT"]
@@ -1085,14 +1088,14 @@ def fixed_field_spec(
                 dtype="bool",
                 shape=(m, self.metadata.num_filters),
                 dimensions=["variants", "filters"],
-                chunks=(variants_chunk_size, self.metadata.num_filters),
+                chunks=(schema_instance.variants_chunk_size, self.metadata.num_filters),
             ),
             fixed_field_spec(
                 name="variant_allele",
                 dtype="O",
                 shape=(m, max_alleles),
                 dimensions=["variants", "alleles"],
-                chunks=(variants_chunk_size, max_alleles),
+                chunks=(schema_instance.variants_chunk_size, max_alleles),
             ),
             fixed_field_spec(
                 name="variant_id",
@@ -1127,7 +1130,10 @@ def fixed_field_spec(
         if gt_field is not None and n > 0:
             ploidy = max(gt_field.summary.max_number - 1, 1)
             shape = [m, n]
-            chunks = [variants_chunk_size, samples_chunk_size]
+            chunks = [
+                schema_instance.variants_chunk_size,
+                schema_instance.samples_chunk_size,
+            ]
             dimensions = ["variants", "samples"]
             array_specs.append(
                 schema.ZarrArraySpec.new(
@@ -1169,15 +1175,8 @@ def fixed_field_spec(
         if local_alleles:
             array_specs = convert_local_allele_field_types(array_specs)
 
-        return schema.VcfZarrSchema(
-            format_version=schema.ZARR_SCHEMA_FORMAT_VERSION,
-            samples_chunk_size=samples_chunk_size,
-            variants_chunk_size=variants_chunk_size,
-            fields=array_specs,
-            samples=self.metadata.samples,
-            contigs=self.metadata.contigs,
-            filters=self.metadata.filters,
-        )
+        schema_instance.fields = array_specs
+        return schema_instance
 
 
 @dataclasses.dataclass

pyproject.toml

@@ -82,9 +82,9 @@ line-length = 88
 indent-width = 4
 
 [tool.ruff.lint]
-select = ["E", "F", "B", "W", "I", "N", "UP", "A", "RUF", "PT"]
+select = ["E", "F", "B", "W", "I", "N", "UP", "A", "PT"]
 #Allow uppercase names for e.g. call_AD
-ignore = ["N806", "N802", "A001", "A002"]
+ignore = ["N806", "N802", "A001", "A002", "RUF"]
 
 fixable = ["ALL"]
 unfixable = []