test: cover annotation lookup helpers

.github/workflows/unit-tests.yml

@@ -0,0 +1,18 @@
+name: Unit tests
+
+on:
+  pull_request:
+    branches: [master]
+
+jobs:
+  unit-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+      - name: Install test dependencies
+        run: python -m pip install pytest gffutils pybedtools pandas numpy future tqdm
+      - name: Run unit tests
+        run: python -m pytest -q unit_tests

annotator/annotation_functions.py

@@ -132,14 +132,12 @@ def get_keys(gtf_format):
 
 def get_gene_to_transcript_dict(db, gene_id_key, transcript_id_key):
     """
-    Returns a gene: transcript dictionary. 
-
-    gene:[transcript1, transcript2]
-
-    :param db: gffutils.FeatureDB
-    :param gene_id_key: string
-    :param transcript_id_key: string
-    :return: 
+    Map each gene ID to the transcript IDs on its transcript features.
+
+    :param db: gffutils.FeatureDB containing transcript features
+    :param gene_id_key: attribute name containing gene IDs
+    :param transcript_id_key: attribute name containing transcript IDs
+    :return: defaultdict(list) mapping gene IDs to lists of transcript IDs
     """
 
     genes_dict = defaultdict(list)
@@ -234,18 +232,12 @@ def get_all_cds_dict(db, cds_key):
 
 def get_all_exons_dict(db, exon_key, transcript_id_key):
     """
-    Returns dictionary of exons as transcript_id:{
-        [
-            {'start':START, 'end':END},
-            {'start':START, 'end':END},
-            ...
-        ]
-    }.
-
-    :param db: gffutils.FeatureDB
-    :param transcript_id_key: string
-    :param exon_key: string
-    :return: 
+    Collect exon coordinates, chromosome, and strand by transcript ID.
+
+    :param db: gffutils.FeatureDB containing exon features
+    :param exon_key: feature type to treat as an exon
+    :param transcript_id_key: attribute name containing transcript IDs
+    :return: defaultdict(list) mapping transcript IDs to exon dictionaries
     """
 
     exons_dict = defaultdict(list)
@@ -264,13 +256,12 @@ def get_all_exons_dict(db, exon_key, transcript_id_key):
 
 def get_all_transcripts_dict(db, transcript_key, transcript_id_key):
     """
-    Returns dictionary of transcript_id:{'start':START, 'end':END}.
-
-    :param db: gffutils.FeatureDB
-    :param transcript_key: string
-    :param transcript_id_key: string
-    :return: chr19_transcripts_dict: defaultdict(dict)
-        hash of transcripts and their start/end coordinates
+    Collect the start and end coordinates of each transcript ID.
+
+    :param db: gffutils.FeatureDB containing transcript features
+    :param transcript_key: feature type to treat as a transcript
+    :param transcript_id_key: attribute name containing transcript IDs
+    :return: defaultdict(dict) mapping transcript IDs to start/end dictionaries
     """
     transcripts_dict = defaultdict(dict)
     for transcript_feature in db.features_of_type(transcript_key):

unit_tests/test_annotation_functions.py

@@ -0,0 +1,121 @@
+"""Unit tests for the public annotation lookup helpers."""
+
+from types import SimpleNamespace
+
+from annotator.annotation_functions import (
+    gene_id_to_name,
+    get_all_exons_dict,
+    get_all_transcripts_dict,
+    get_gene_to_transcript_dict,
+)
+
+
+class FeatureDBStub:
+    def __init__(self, features):
+        self.features = features
+
+    def features_of_type(self, feature_type):
+        return iter(self.features.get(feature_type, ()))
+
+
+def feature(attributes, start=1, end=10, seqid="chr1", strand="+"):
+    return SimpleNamespace(
+        attributes=attributes,
+        start=start,
+        end=end,
+        seqid=seqid,
+        strand=strand,
+    )
+
+
+def test_gene_to_transcript_dict_groups_transcripts_by_gene():
+    db = FeatureDBStub(
+        {
+            "transcript": [
+                feature({"gene_id": ["gene-a"], "transcript_id": ["tx-1"]}),
+                feature({"gene_id": ["gene-a"], "transcript_id": ["tx-2"]}),
+                feature({"gene_id": ["gene-b"], "transcript_id": ["tx-3"]}),
+            ],
+            "gene": [feature({"gene_id": ["unrelated"]})],
+        }
+    )
+
+    assert get_gene_to_transcript_dict(db, "gene_id", "transcript_id") == {
+        "gene-a": ["tx-1", "tx-2"],
+        "gene-b": ["tx-3"],
+    }
+
+
+def test_gene_to_transcript_dict_uses_supplied_attribute_keys():
+    db = FeatureDBStub(
+        {
+            "transcript": [
+                feature({"gene": ["gene-a", "gene-b"], "isoform": ["tx-1", "tx-2"]})
+            ]
+        }
+    )
+
+    assert get_gene_to_transcript_dict(db, "gene", "isoform") == {
+        "gene-a": ["tx-1", "tx-2"],
+        "gene-b": ["tx-1", "tx-2"],
+    }
+
+
+def test_gene_to_transcript_dict_returns_empty_mapping_without_transcripts():
+    db = FeatureDBStub({"gene": [feature({"gene_id": ["gene-a"]})]})
+
+    assert get_gene_to_transcript_dict(db, "gene_id", "transcript_id") == {}
+
+
+def test_all_transcripts_dict_preserves_coordinates_for_each_transcript_id():
+    db = FeatureDBStub(
+        {
+            "mRNA": [
+                feature({"ID": ["tx-1", "tx-2"]}, start=10, end=30),
+                feature({"ID": ["tx-3"]}, start=40, end=70),
+            ],
+            "transcript": [feature({"ID": ["unrelated"]})],
+        }
+    )
+
+    assert get_all_transcripts_dict(db, "mRNA", "ID") == {
+        "tx-1": {"start": 10, "end": 30},
+        "tx-2": {"start": 10, "end": 30},
+        "tx-3": {"start": 40, "end": 70},
+    }
+
+
+def test_all_exons_dict_keeps_each_exon_and_its_strand():
+    db = FeatureDBStub(
+        {
+            "exon": [
+                feature({"transcript_id": ["tx-1", "tx-2"]}, 10, 20, "chr2", "-"),
+                feature({"transcript_id": ["tx-1"]}, 30, 40, "chr2", "-"),
+            ]
+        }
+    )
+
+    assert get_all_exons_dict(db, "exon", "transcript_id") == {
+        "tx-1": [
+            {"chrom": "chr2", "start": 10, "end": 20, "strand": "-"},
+            {"chrom": "chr2", "start": 30, "end": 40, "strand": "-"},
+        ],
+        "tx-2": [{"chrom": "chr2", "start": 10, "end": 20, "strand": "-"}],
+    }
+
+
+def test_gene_id_to_name_accepts_list_and_scalar_gene_ids():
+    db = FeatureDBStub(
+        {
+            "gene": [
+                feature({"gene_id": ["gene-a"], "symbol": ["Alpha"]}),
+                feature({"gene_id": "gene-b", "symbol": ["Beta"]}),
+            ],
+            "transcript": [feature({"gene_id": ["unrelated"], "symbol": ["Other"]})],
+        }
+    )
+
+    assert gene_id_to_name(db, "symbol") == {
+        "gene-a": "Alpha",
+        "gene-b": "Beta",
+    }