feat: get basic tests working

Author: ameynert

divref/divref/haplotype.py

@@ -26,7 +26,7 @@ def to_hashable_items(d: dict[str, _V]) -> tuple[tuple[str, _V], ...]:
     return tuple(sorted(d.items()))
 
 
-def get_haplo_sequence(context_size: int, variants: Any) -> Any:
+def get_haplo_sequence(context_size: int, variants: Any, reference_genome: str = "GRCh38") -> Any:
     """
     Construct a haplotype sequence string with flanking genomic context.
 
@@ -38,6 +38,7 @@ def get_haplo_sequence(context_size: int, variants: Any) -> Any:
         context_size: Number of reference bases to include flanking each end.
         variants: Hail array expression of variant structs with locus and alleles fields.
             Must contain at least one variant.
+        reference_genome: Name of the reference genome. Defaults to "GRCh38".
 
     Returns:
         Hail string expression representing the full haplotype sequence.
@@ -60,7 +61,7 @@ def get_haplo_sequence(context_size: int, variants: Any) -> Any:
         min_pos,
         before=context_size,
         after=(max_pos - min_pos + max_variant_size + context_size - 1),
-        reference_genome="GRCh38",
+        reference_genome=reference_genome,
     )
 
     # (min_pos - index_translation) equals context_size, mapping locus positions to string indices
@@ -117,7 +118,7 @@ def split_haplotypes(ht: Any, window_size: int) -> Any:
         Hail table with haplotypes exploded into sub-haplotypes by window.
     """
     breakpoints = hl.range(1, hl.len(ht.variants)).filter(
-        lambda i: (i == 0) | (variant_distance(ht.variants[i - 1], ht.variants[i]) >= window_size)
+        lambda i: variant_distance(ht.variants[i - 1], ht.variants[i]) >= window_size
     )
 
     def get_range(i: Any) -> Any:

divref/tests/conftest.py

@@ -1,6 +1,8 @@
 """Shared pytest fixtures for the divref test suite."""
 
 from collections.abc import Generator
+from pathlib import Path
+from typing import Any
 
 import hail as hl
 import pytest
@@ -18,3 +20,20 @@ def hail_context() -> Generator[None, None, None]:
     hl.init(quiet=True)
     yield
     hl.stop()
+
+
+@pytest.fixture
+def datadir() -> Path:
+    return Path(__file__).parent / "data"
+
+
+@pytest.fixture
+def hail_reference_genome() -> Any:
+    """A small custom reference genome for use in testing."""
+    contigs: list[str] = ["chr1"]
+    lengths: dict[str, int] = {"chr1": 1000}
+
+    reference_genome = hl.ReferenceGenome(
+        "test_chr1", contigs, lengths, x_contigs=[], y_contigs=[], mt_contigs=[]
+    )
+    return reference_genome

divref/tests/data/test.fa

@@ -0,0 +1,11 @@
+> chr1
+AGCTGGTGTTCGGTTCGGTAACGGAGAATCTGTGGGGCTATGTCACTAATACTTTCGAAACGCCCCGTACCGATGCTGAACAAGTCGATGCAGGCTCCCG
+TCTTTGAATAGGGGTAAACATACAAGTCGATAGAAGATGGGTAGGGGCCTCCAATTCATCCAACACTCTACGCCTTCTCCAAGAGCTAGTAGGGCACCCT
+GCAGTTGGAAAGGGAACTATTTCGTAGGGCGAGCCCATACCGTCTCTCTTGCGGAAGACTTAACACGATAGGAAGCTGGAATAGTTTCGAACGATGGTTA
+TTAATCCTAATAACGGAACGCTGTCTGGAGGATGAGTGTGACGGAGTGTAACTCGATGAGTTACCCGCTAATCGAACTGGGCGAGAGATCCCAGCGCTGA
+TGCACTCGATCCCGAGGCCTGACCCGACATATCAGCTCAGACTAGAGCGGGGCTGTTGACGTTTGGGGTTGAAAAAATCTATTGTACCAATCGGCTTCAA
+CGTGCTCCACGGCTGGCGCCTGAGGAGGGGCCCACACCGAGGAAGTAGACTGTTGCACGTTGGCGATGGCGGTAGCTAACTAAGTCGCCTGCCACAACAA
+CAGTATCAAAGCCGTATAAAGGGAACATCCACACTTTAGTGAATCGAAGCGCGGCATCAGAATTTCCTTTTGGATACCTGATACAAAGCCCATCGTGGTC
+CTTAGACTTCGTGCACATACAGCTGCACCGCACGCATGTGGAATTAGAGGCGAAGTACGATTCCTAGACCGACGTACGATACAACTATGTGGATGTGACG
+AGCTTCTTTTATATGCTTCGCCCGCCGGACCGGCCTCGCGATGGCGTAGCTGCGCATAAGCAAATGACAATTAACCACTGTGTACTCGTTATAACATCTG
+GCAGTTAAAGTCGGGAGAATAGGAGCCGCAATACACAGTTTACCGCATCTAGACCTAACTGAGATACTGCCATAGACGACTAGCCATCCCTCTGGCTCTT

divref/tests/data/test.fa.fai

@@ -0,0 +1 @@
+chr1	1000	7	100	101

divref/tests/test_haplotype.py

@@ -1,22 +1,70 @@
 """Tests for shared Hail utilities in haplotype.py."""
 
+from pathlib import Path
 from typing import Any
 
+import hail as hl
 import pytest
 
 from divref.haplotype import get_haplo_sequence
 from divref.haplotype import split_haplotypes
 from divref.haplotype import to_hashable_items
 from divref.haplotype import variant_distance
 
-hl = pytest.importorskip("hail")
+# ---------------------------------------------------------------------------
+# Helper functions
+# ---------------------------------------------------------------------------
+
+
+def _make_variant(position: int, ref: str, alt: str) -> Any:
+    return hl.Struct(locus=hl.Struct(contig="chr1", position=position), alleles=[ref, alt])
+
+
+def _make_haplotype_table(variant_positions: list[tuple[int, str, str]]) -> Any:
+    variant_type = hl.tstruct(locus=hl.tstruct(position=hl.tint32), alleles=hl.tarray(hl.tstr))
+    row_type = hl.tstruct(
+        variants=hl.tarray(variant_type),
+        haplotype=hl.tarray(hl.tstr),
+        gnomad_freqs=hl.tarray(hl.tfloat64),
+    )
+    variants = [
+        {"locus": {"position": pos}, "alleles": [ref, alt]} for pos, ref, alt in variant_positions
+    ]
+    return hl.Table.parallelize(
+        [
+            {
+                "variants": variants,
+                "haplotype": [str(i) for i in range(len(variants))],
+                "gnomad_freqs": [0.1] * len(variants),
+            }
+        ],
+        schema=row_type,
+    )
 
 
 # ---------------------------------------------------------------------------
 # get_haplo_sequence
 # ---------------------------------------------------------------------------
 
 
+def test_get_haplo_sequence_single(
+    datadir: Path,
+    hail_reference_genome: Any,
+    hail_context: None,  # noqa: ARG001
+) -> None:
+    """get_haplo_sequence should return the correct haplotype sequence."""
+    test_fasta: Path = datadir / "test.fa"
+    test_fai: Path = datadir / "test.fa.fai"
+
+    hail_reference_genome.add_sequence(str(test_fasta), str(test_fai))
+
+    variant = _make_variant(position=100, ref="A", alt="C")
+    x = get_haplo_sequence(
+        context_size=2, variants=[variant], reference_genome=hail_reference_genome.name
+    )
+    print(hl.str(x))
+
+
 def test_get_haplo_sequence_empty_list_raises() -> None:
     """get_haplo_sequence should raise ValueError when given an empty list."""
     with pytest.raises(ValueError, match="at least one variant"):
@@ -51,10 +99,6 @@ def test_to_hashable_items_sorted_by_key() -> None:
 # ---------------------------------------------------------------------------
 
 
-def _make_variant(position: int, ref: str, alt: str) -> Any:
-    return hl.Struct(locus=hl.Struct(position=position), alleles=[ref, alt])
-
-
 def test_variant_distance_adjacent_snps(hail_context: None) -> None:  # noqa: ARG001
     # SNP at 100, next SNP at 101: distance = 101 - 100 - len("A") = 0
     assert (
@@ -81,28 +125,6 @@ def test_variant_distance_deletion_closes_gap(hail_context: None) -> None:  # no
 # ---------------------------------------------------------------------------
 
 
-def _make_haplotype_table(variant_positions: list[tuple[int, str, str]]) -> Any:
-    variant_type = hl.tstruct(locus=hl.tstruct(position=hl.tint32), alleles=hl.tarray(hl.tstr))
-    row_type = hl.tstruct(
-        variants=hl.tarray(variant_type),
-        haplotype=hl.tarray(hl.tstr),
-        gnomad_freqs=hl.tarray(hl.tfloat64),
-    )
-    variants = [
-        {"locus": {"position": pos}, "alleles": [ref, alt]} for pos, ref, alt in variant_positions
-    ]
-    return hl.Table.parallelize(
-        [
-            {
-                "variants": variants,
-                "haplotype": [str(i) for i in range(len(variants))],
-                "gnomad_freqs": [0.1] * len(variants),
-            }
-        ],
-        schema=row_type,
-    )
-
-
 def test_split_haplotypes_no_split_needed(hail_context: None) -> None:  # noqa: ARG001
     # All variants within window_size=200; haplotype is kept intact as one row
     ht = _make_haplotype_table([(100, "A", "T"), (150, "C", "G"), (190, "G", "A")])