Add JsonDataclass for asdict/asjson

Author: jeromekelleher

bio2zarr/core.py

@@ -1,6 +1,7 @@
 import concurrent.futures as cf
 import contextlib
 import dataclasses
+import json
 import logging
 import multiprocessing
 import os
@@ -277,3 +278,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self._update_progress()
         self.progress_bar.close()
         return False
+
+
+class JsonDataclass:
+    def asdict(self):
+        return dataclasses.asdict(self)
+
+    def asjson(self):
+        return json.dumps(self.asdict(), indent=4)

bio2zarr/vcf.py

@@ -50,7 +50,7 @@ def display_size(n):
 
 
 @dataclasses.dataclass
-class VcfFieldSummary:
+class VcfFieldSummary(core.JsonDataclass):
     num_chunks: int = 0
     compressed_size: int = 0
     uncompressed_size: int = 0
@@ -67,9 +67,6 @@ def update(self, other):
         self.min_value = min(self.min_value, other.min_value)
         self.max_value = max(self.max_value, other.max_value)
 
-    def asdict(self):
-        return dataclasses.asdict(self)
-
     @staticmethod
     def fromdict(d):
         return VcfFieldSummary(**d)
@@ -168,7 +165,7 @@ class Filter:
 
 
 @dataclasses.dataclass
-class IcfMetadata:
+class IcfMetadata(core.JsonDataclass):
     samples: list
     contigs: list
     filters: list
@@ -226,12 +223,6 @@ def fromdict(d):
         d["contigs"] = [Contig(**cd) for cd in d["contigs"]]
         return IcfMetadata(**d)
 
-    def asdict(self):
-        return dataclasses.asdict(self)
-
-    def asjson(self):
-        return json.dumps(self.asdict(), indent=4)
-
 
 def fixed_vcf_field_definitions():
     def make_field_def(name, vcf_type, vcf_number):
@@ -933,17 +924,11 @@ def num_fields(self):
 
 
 @dataclasses.dataclass
-class IcfPartitionMetadata:
+class IcfPartitionMetadata(core.JsonDataclass):
     num_records: int
     last_position: int
     field_summaries: dict
 
-    def asdict(self):
-        return dataclasses.asdict(self)
-
-    def asjson(self):
-        return json.dumps(self.asdict(), indent=4)
-
     @staticmethod
     def fromdict(d):
         md = IcfPartitionMetadata(**d)
@@ -987,17 +972,11 @@ def check_field_clobbering(icf_metadata):
 
 
 @dataclasses.dataclass
-class IcfWriteSummary:
+class IcfWriteSummary(core.JsonDataclass):
     num_partitions: int
     num_samples: int
     num_variants: int
 
-    def asdict(self):
-        return dataclasses.asdict(self)
-
-    def asjson(self):
-        return json.dumps(self.asdict(), indent=4)
-
 
 class IntermediateColumnarFormatWriter:
     def __init__(self, path):
@@ -1409,7 +1388,7 @@ def variant_chunk_nbytes(self):
 
 
 @dataclasses.dataclass
-class VcfZarrSchema:
+class VcfZarrSchema(core.JsonDataclass):
     format_version: str
     samples_chunk_size: int
     variants_chunk_size: int
@@ -1421,12 +1400,6 @@ class VcfZarrSchema:
     def field_map(self):
         return {field.name: field for field in self.fields}
 
-    def asdict(self):
-        return dataclasses.asdict(self)
-
-    def asjson(self):
-        return json.dumps(self.asdict(), indent=4)
-
     @staticmethod
     def fromdict(d):
         if d["format_version"] != ZARR_SCHEMA_FORMAT_VERSION:
@@ -1645,17 +1618,14 @@ def generate_partitions(num_records, chunk_size, num_partitions, max_chunks=None
 
 
 @dataclasses.dataclass
-class VcfZarrWriterMetadata:
+class VcfZarrWriterMetadata(core.JsonDataclass):
     format_version: str
     icf_path: str
     schema: VcfZarrSchema
     dimension_separator: str
     partitions: list
     provenance: dict
 
-    def asdict(self):
-        return dataclasses.asdict(self)
-
     @staticmethod
     def fromdict(d):
         if d["format_version"] != VZW_METADATA_FORMAT_VERSION:
@@ -1670,19 +1640,13 @@ def fromdict(d):
 
 
 @dataclasses.dataclass
-class VcfZarrWriteSummary:
+class VcfZarrWriteSummary(core.JsonDataclass):
     num_partitions: int
     num_samples: int
     num_variants: int
     num_chunks: int
     max_encoding_memory: str
 
-    def asdict(self):
-        return dataclasses.asdict(self)
-
-    def asjson(self):
-        return json.dumps(self.asdict(), indent=4)
-
 
 class VcfZarrWriter:
     def __init__(self, path):