Minimal changes to get things working again

Author: jeromekelleher

bio2zarr/vcf.py

@@ -25,7 +25,7 @@
 
 from . import core
 from . import provenance
-from .vcf_partition import partition_into_regions, region_filter
+from . import vcf_utils
 
 logger = logging.getLogger(__name__)
 
@@ -229,8 +229,11 @@ def scan_vcfs(paths, show_progress, target_num_partitions):
                 raise ValueError("Incompatible VCF chunks")
         vcf_metadata.num_records += vcf.num_records
 
-        # https://github.com/pystatgen/sgkit/issues/1200
-        regions = partition_into_regions(path, num_parts=target_num_partitions)
+        # TODO: Move all our usage of the VCF class behind the IndexedVCF
+        # so that we open the VCF once, and we explicitly set the index.
+        # Otherwise cyvcf2 will do things behind our backs.
+        indexed_vcf = vcf_utils.IndexedVcf(path)
+        regions = indexed_vcf.partition_into_regions(num_parts=target_num_partitions)
         for region in regions:
             partitions.append(
                 # Requires cyvcf2>=0.30.27
@@ -241,8 +244,7 @@ def scan_vcfs(paths, show_progress, target_num_partitions):
             )
     # TODO figure out if this is safe when we have multiple chrs
     # in the file
-    # FIXME this isn't working because region strings don't sort correctly
-    partitions.sort(key=lambda x: x.region)
+    partitions.sort(key=lambda x: (x.region.contig, x.region.start))
     vcf_metadata.partitions = partitions
     return vcf_metadata, header
 
@@ -805,14 +807,20 @@ def convert_partition(
                     format_fields.append(field)
 
         def variants():
-            with warnings.catch_warnings():
-                # TODO cyvcf2 emits a warning for empty regions; either make the
-                # warning more specific, or remove the need for querying empty
-                # regions.
-                # FIXME this also absorbs any warnings emitted within the loop,
-                # so definitely need to do this a different way.
-                warnings.simplefilter("ignore")
-                for var in region_filter(vcf(partition.region), partition.region):
+            # with warnings.catch_warnings():
+            #     # TODO cyvcf2 emits a warning for empty regions; either make the
+            #     # warning more specific, or remove the need for querying empty
+            #     # regions.
+            #     # FIXME this also absorbs any warnings emitted within the loop,
+            #     # so definitely need to do this a different way.
+            #     warnings.simplefilter("ignore")
+            #     for var in region_filter(vcf(partition.region), partition.region):
+            #         yield var
+
+            # TODO move this into the IndexedVCF class
+            start = 1 if partition.region.start is None else partition.region.start
+            for var in vcf(str(partition.region)):
+                if var.POS >= start:
                     yield var
 
         # FIXME it looks like this is actually a bit pointless now that we

bio2zarr/vcf_utils.py

@@ -52,11 +52,18 @@ def region_string(contig: str, start: int, end: Optional[int] = None) -> str:
     else:
         return f"{contig}:{start}-"
 
+
 @dataclass
 class Region:
     contig: str
     start: Optional[int] = None
-    end: Optional[int]=None
+    end: Optional[int] = None
+
+    def __post_init__(self):
+        if self.start is not None:
+            self.start = int(self.start)
+        if self.end is not None:
+            self.end = int(self.end)
 
     def __str__(self):
         s = f"{self.contig}"
@@ -68,6 +75,7 @@ def __str__(self):
 
     # TODO add "parse" class method
 
+
 def get_tabix_path(
     vcf_path: PathType, storage_options: Optional[Dict[str, str]] = None
 ) -> Optional[str]:
@@ -530,18 +538,19 @@ def read_tabix(
         )
 
 
-
 class IndexedVcf:
-    def __init__(self, path, index_path=None):
-        # for h in vcf.header_iter():
-        #     print(h)
-        # if index_path is None:
-        #     index_path = get_tabix_path(vcf_path, storage_options=storage_options)
-        #     if index_path is None:
-        #         index_path = get_csi_path(vcf_path, storage_options=storage_options)
-        #         if index_path is None:
-        #             raise ValueError("Cannot find .tbi or .csi file.")
-        self.vcf_path = path
+    def __init__(self, vcf_path, index_path=None):
+        vcf_path = pathlib.Path(vcf_path)
+        if index_path is None:
+            index_path = vcf_path.with_suffix(vcf_path.suffix + ".tbi")
+            if not index_path.exists():
+                index_path = vcf_path.with_suffix(vcf_path.suffix + ".csi")
+                if not index_path.exists():
+                    raise ValueError("Cannot find .tbi or .csi file.")
+        else:
+            index_path = pathlib.Path(index_path)
+
+        self.vcf_path = vcf_path
         self.index_path = index_path
         self.file_type = None
         self.index_type = None
@@ -561,7 +570,7 @@ def __init__(self, path, index_path=None):
                 self.file_type = "vcf"
                 self.sequence_names = self.index.parse_vcf_aux()
             else:
-                with contextlib.closing(cyvcf2.VCF(path)) as vcf:
+                with contextlib.closing(cyvcf2.VCF(vcf_path)) as vcf:
                     self.sequence_names = vcf.seqnames
         else:
             self.sequence_names = self.index.sequence_names