Good implementation of no-indexes for IndexedVcf

Author: jeromekelleher

bio2zarr/vcf_utils.py

@@ -89,9 +89,7 @@ class Region:
     end: Optional[int] = None
 
     def __post_init__(self):
-        if self.contig is None:
-            return
-
+        assert self.contig is not None
         if self.start is not None:
             self.start = int(self.start)
             assert self.start > 0
@@ -399,6 +397,9 @@ class VcfIndexType(Enum):
 class IndexedVcf(contextlib.AbstractContextManager):
     def __init__(self, vcf_path, index_path=None):
         self.vcf = None
+        self.file_type = None
+        self.index_type = None
+
         vcf_path = pathlib.Path(vcf_path)
         if not vcf_path.exists():
             raise FileNotFoundError(vcf_path)
@@ -411,27 +412,28 @@ def __init__(self, vcf_path, index_path=None):
                     vcf_path.suffix + VcfIndexType.CSI.value
                 )
                 if not index_path.exists():
-                    # Use this as a proxy for "no index"
-                    index_path = vcf_path
+                    # No supported index found
+                    index_path = None
         else:
             index_path = pathlib.Path(index_path)
+            if not index_path.exists():
+                raise FileNotFoundError(
+                    f"Specified index path {index_path} does not exist"
+                )
 
         self.vcf_path = vcf_path
         self.index_path = index_path
-        self.file_type = None
-        self.index_type = None
-
-        if index_path.suffix == VcfIndexType.CSI.value:
-            self.index_type = VcfIndexType.CSI
-        elif index_path.suffix == VcfIndexType.TABIX.value:
-            self.index_type = VcfIndexType.TABIX
-            self.file_type = VcfFileType.VCF
-        # else:
-
-        #     raise ValueError("Only .tbi or .csi indexes are supported.")
+        if index_path is not None:
+            if index_path.suffix == VcfIndexType.CSI.value:
+                self.index_type = VcfIndexType.CSI
+            elif index_path.suffix == VcfIndexType.TABIX.value:
+                self.index_type = VcfIndexType.TABIX
+                self.file_type = VcfFileType.VCF
+            else:
+                raise ValueError("Only .tbi or .csi indexes are supported.")
 
         self.vcf = cyvcf2.VCF(vcf_path)
-        if self.index_type is not None:
+        if self.index_path is not None:
             self.vcf.set_index(str(self.index_path))
 
         logger.debug(f"Loaded {vcf_path} with index {self.index_path}")
@@ -449,7 +451,15 @@ def __init__(self, vcf_path, index_path=None):
                 self.sequence_names = self.vcf.seqnames
         elif self.index_type == VcfIndexType.TABIX:
             self.index = read_tabix(self.index_path)
+            self.file_type = VcfFileType.VCF
             self.sequence_names = self.index.sequence_names
+        else:
+            assert self.index is None
+            var = next(self.vcf)
+            self.sequence_names = [var.CHROM]
+            self.vcf.close()
+            # There doesn't seem to be a way to reset the iterator
+            self.vcf = cyvcf2.VCF(vcf_path)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if self.vcf is not None:
@@ -459,7 +469,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     def contig_record_counts(self):
         if self.index is None:
-            return {None: np.inf}
+            return {self.sequence_names[0]: np.inf}
         d = dict(zip(self.sequence_names, self.index.record_counts))
         if self.file_type == VcfFileType.BCF:
             d = {k: v for k, v in d.items() if v > 0}
@@ -468,10 +478,15 @@ def contig_record_counts(self):
     def count_variants(self, region):
         return sum(1 for _ in self.variants(region))
 
-    def variants(self, region):
+    def variants(self, region=None):
         if self.index is None:
-            assert region.contig is None
-            yield from self.vcf
+            contig = self.sequence_names[0]
+            if region is not None:
+                assert region.contig == contig
+            for var in self.vcf:
+                if var.CHROM != contig:
+                    raise ValueError("Multi-contig VCFs must be indexed")
+                yield var
         else:
             start = 1 if region.start is None else region.start
             for var in self.vcf(str(region)):
@@ -498,9 +513,6 @@ def partition_into_regions(
         num_parts: Optional[int] = None,
         target_part_size: Union[None, int, str] = None,
     ):
-        if self.index is None:
-            return [Region()]
-
         if num_parts is None and target_part_size is None:
             raise ValueError("One of num_parts or target_part_size must be specified")
 
@@ -520,6 +532,9 @@ def partition_into_regions(
             if target_part_size_bytes < 1:
                 raise ValueError("target_part_size must be positive")
 
+        if self.index is None:
+            return [Region(self.sequence_names[0])]
+
         # Calculate the desired part file boundaries
         file_length = os.stat(self.vcf_path).st_size
         if num_parts is not None:

tests/test_vcf_utils.py

@@ -1,4 +1,5 @@
 import pathlib
+import shutil
 
 import numpy as np
 import pytest
@@ -159,6 +160,27 @@ def test_partition_into_n_parts(self, index_file, total_records, num_parts):
         assert np.sum(part_variant_counts) == total_records
         assert_part_counts_non_zero(part_variant_counts, index_file)
 
+    @pytest.mark.parametrize(
+        ("vcf_file", "total_records"),
+        [
+            ("1kg_2020_chrM.vcf.gz", 23),
+            ("1kg_2020_chr20_annotations.bcf", 21),
+        ],
+    )
+    @pytest.mark.parametrize("num_parts", [1, 2, 3])
+    def test_partition_into_n_parts_unindexed(
+        self, tmp_path, vcf_file, total_records, num_parts
+    ):
+        copy_path = tmp_path / vcf_file
+        shutil.copyfile(data_path / vcf_file, copy_path)
+        indexed_vcf = vcf_utils.IndexedVcf(copy_path)
+        regions = list(indexed_vcf.partition_into_regions(num_parts=num_parts))
+        assert len(regions) == 1
+        part_variant_counts = np.array(
+            [indexed_vcf.count_variants(region) for region in regions]
+        )
+        assert np.sum(part_variant_counts) == total_records
+
     def test_tabix_multi_chrom_bug(self):
         indexed_vcf = self.get_instance("multi_contig.vcf.gz.tbi")
         regions = list(indexed_vcf.partition_into_regions(num_parts=10))
@@ -213,11 +235,54 @@ def test_partition_invalid_arguments(self):
         with pytest.raises(ValueError, match=r"target_part_size must be positive"):
             indexed_vcf.partition_into_regions(target_part_size=0)
 
-    def test_bad_index(self):
-        with pytest.raises(
-            ValueError, match=r"Only .tbi or .csi indexes are supported."
-        ):
-            vcf_utils.IndexedVcf(data_path / "sample.vcf.gz", "y")
+    @pytest.mark.parametrize("path", ["y", data_path / "xxx", "/x/y.csi"])
+    def test_missing_index_file(self, path):
+        with pytest.raises(FileNotFoundError, match="Specified index path"):
+            vcf_utils.IndexedVcf(data_path / "sample.vcf.gz", path)
+
+    def test_bad_index_format(self):
+        vcf_file = data_path / "sample.vcf.gz"
+        with pytest.raises(ValueError, match="Only .tbi or .csi indexes"):
+            vcf_utils.IndexedVcf(vcf_file, vcf_file)
+
+    @pytest.mark.parametrize(
+        "filename",
+        [
+            "1kg_2020_chrM.vcf.gz",
+            "1kg_2020_chrM.bcf",
+            "1kg_2020_chr20_annotations.bcf",
+            "chr_m_indels.vcf.gz",
+            "NA12878.prod.chr20snippet.g.vcf.gz",
+        ],
+    )
+    def test_unindexed_single_contig(self, tmp_path, filename):
+        f1 = vcf_utils.IndexedVcf(data_path / filename)
+        assert f1.index is not None
+        copy_path = tmp_path / filename
+        shutil.copyfile(data_path / filename, copy_path)
+        f2 = vcf_utils.IndexedVcf(copy_path)
+        assert f2.index is None
+        crc1 = f1.contig_record_counts()
+        assert len(crc1) == 1
+        contig = next(iter(crc1.keys()))
+        assert f2.contig_record_counts() == {contig: np.inf}
+        region = vcf_utils.Region(contig)
+        # The full variants returned by cyvcf2 don't compare for equality,
+        # so just check the chrom/pos values
+        v1 = [(v.CHROM, v.POS) for v in f1.variants(region)]
+        v2 = [(v.CHROM, v.POS) for v in f2.variants(region)]
+        assert v1 == v2
+
+    @pytest.mark.parametrize(
+        "filename",
+        ["sample.vcf.gz", "sample.bcf", "multi_contig.vcf.gz"],
+    )
+    def test_unindexed_multi_contig(self, tmp_path, filename):
+        copy_path = tmp_path / filename
+        shutil.copyfile(data_path / filename, copy_path)
+        f = vcf_utils.IndexedVcf(copy_path)
+        with pytest.raises(ValueError, match="Multi-contig VCFs must be indexed"):
+            list(f.variants())
 
 
 class TestCsiIndex: