Change "spec" to schema, and minor tweak on format

Author: jeromekelleher

bio2zarr/cli.py

@@ -56,10 +56,9 @@ def inspect(if_path, verbose):
 
 @click.command
 @click.argument("if_path", type=click.Path())
-# @click.argument("specfile", type=click.Path())
-def genspec(if_path):
+def mkschema(if_path):
     stream = click.get_text_stream("stdout")
-    vcf.generate_spec(if_path, stream)
+    vcf.mkschema(if_path, stream)
 
 
 @click.command
@@ -105,7 +104,7 @@ def vcf2zarr():
 
 vcf2zarr.add_command(explode)
 vcf2zarr.add_command(inspect)
-vcf2zarr.add_command(genspec)
+vcf2zarr.add_command(mkschema)
 vcf2zarr.add_command(to_zarr)
 vcf2zarr.add_command(convert_vcf)
 vcf2zarr.add_command(validate)

bio2zarr/vcf.py

@@ -887,7 +887,6 @@ def inspect(if_path):
 
 @dataclasses.dataclass
 class ZarrColumnSpec:
-    # TODO change to "variable_name"
     name: str
     dtype: str
     shape: tuple
@@ -898,6 +897,11 @@ class ZarrColumnSpec:
     compressor: dict
     # TODO add filters
 
+    def __post_init__(self):
+        self.shape = tuple(self.shape)
+        self.chunks = tuple(self.chunks)
+        self.dimensions = tuple(self.dimensions)
+
 
 @dataclasses.dataclass
 class ZarrConversionSpec:
@@ -908,17 +912,24 @@ class ZarrConversionSpec:
     contig_id: list
     contig_length: list
     filter_id: list
-    variables: list
+    columns: dict
 
     def asdict(self):
         return dataclasses.asdict(self)
 
+    def asjson(self):
+        return json.dumps(self.asdict(), indent=4)
+
     @staticmethod
     def fromdict(d):
         ret = ZarrConversionSpec(**d)
-        ret.variables = [ZarrColumnSpec(**cd) for cd in d["variables"]]
+        ret.columns = {key: ZarrColumnSpec(**value) for key,value in d["columns"].items()}
         return ret
 
+    @staticmethod
+    def fromjson(s):
+        return ZarrConversionSpec.fromdict(json.loads(s))
+
     @staticmethod
     def generate(pcvcf, chunk_length=None, chunk_width=None):
         m = pcvcf.num_records
@@ -1070,7 +1081,7 @@ def fixed_field_spec(
         return ZarrConversionSpec(
             chunk_width=chunk_width,
             chunk_length=chunk_length,
-            variables=colspecs,
+            columns={col.name: col for col in colspecs},
             dimensions=["variants", "samples", "ploidy", "alleles", "filters"],
             sample_id=pcvcf.metadata.samples,
             contig_id=pcvcf.metadata.contig_names,
@@ -1261,8 +1272,8 @@ def convert(
         logger.info(f"Create zarr at {write_path}")
         sgvcf = SgvcfZarr(write_path)
         sgvcf.root = zarr.group(store=store, overwrite=True)
-        for variable in conversion_spec.variables[:]:
-            sgvcf.create_array(variable)
+        for column in conversion_spec.columns.values():
+            sgvcf.create_array(column)
 
         progress_config = core.ProgressConfig(
             total=pcvcf.total_uncompressed_bytes,
@@ -1287,7 +1298,7 @@ def convert(
             )
             pwm.submit(sgvcf.encode_filters, pcvcf, conversion_spec.filter_id)
             has_gt = False
-            for variable in conversion_spec.variables[:]:
+            for variable in conversion_spec.columns.values():
                 if variable.vcf_field is not None:
                     # print("Encode", variable.name)
                     # TODO for large columns it's probably worth splitting up
@@ -1309,10 +1320,10 @@ def convert(
         os.rename(write_path, path)
 
 
-def generate_spec(if_path, out):
+def mkschema(if_path, out):
     pcvcf = PickleChunkedVcf.load(if_path)
     spec = ZarrConversionSpec.generate(pcvcf)
-    json.dump(spec.asdict(), out, indent=4)
+    out.write(spec.asjson())
 
 
 def to_zarr(
@@ -1323,8 +1334,7 @@ def to_zarr(
         spec = ZarrConversionSpec.generate(pcvcf)
     else:
         with open(conversion_spec, "r") as f:
-            d = json.load(f)
-            spec = ZarrConversionSpec.fromdict(d)
+            spec = ZarrConversionSpec.fromjson(f.read())
     SgvcfZarr.convert(
         pcvcf,
         zarr_path,

tests/test_cli.py

@@ -33,3 +33,17 @@ def test_inspect(self):
             assert result.stdout == "\n"
             assert len(result.stderr) == 0
             mocked.assert_called_once_with("path")
+
+    def test_mkschema(self):
+        runner = ct.CliRunner(mix_stderr=False)
+        with mock.patch("bio2zarr.vcf.mkschema") as mocked:
+            result = runner.invoke(
+                cli.vcf2zarr, ["mkschema", "path"], catch_exceptions=False
+            )
+            assert result.exit_code == 0
+            assert len(result.stdout) == 0
+            assert len(result.stderr) == 0
+            # TODO figure out how to test that we call it with stdout from
+            # the CliRunner
+            # mocked.assert_called_once_with("path", stdout)
+            mocked.assert_called_once()

tests/test_pcvcf.py

@@ -22,6 +22,15 @@ def pcvcf(self, tmp_path_factory):
         out = tmp_path_factory.mktemp("data") / "example.exploded"
         return vcf.explode([self.data_path], out)
 
+    def test_mkschema(self, tmp_path, pcvcf):
+        schema_file = tmp_path / "schema.json"
+        with open(schema_file, "w") as f:
+            vcf.mkschema(pcvcf.path, f)
+        with open(schema_file, "r") as f:
+            schema1 = vcf.ZarrConversionSpec.fromjson(f.read())
+        schema2 = vcf.ZarrConversionSpec.generate(pcvcf)
+        assert schema1 == schema2
+
     def test_summary_table(self, pcvcf):
         data = pcvcf.summary_table()
         cols = [d["name"] for d in data]
@@ -110,8 +119,7 @@ def schema(self, pcvcf):
         ],
     )
     def test_info_schemas(self, schema, name, dtype, shape):
-        variables = [v for v in schema.variables if v.name == name]
-        v = variables[0]
+        v = schema.columns[name]
         assert v.dtype == dtype
         assert tuple(v.shape) == shape
 

tests/test_vcf_examples.py

@@ -293,7 +293,7 @@ def test_full_pipeline(self, ds, tmp_path, worker_processes):
         )
         schema = tmp_path / "schema.json"
         with open(schema, "w") as f:
-            vcf.generate_spec(exploded, f)
+            vcf.mkschema(exploded, f)
         out = tmp_path / "example.zarr"
         vcf.to_zarr(exploded, out, schema, worker_processes=worker_processes)
         ds2 = sg.load_dataset(out)