feat: add anyvlm client for caf retrieval

close #16

Author: korikuzma

pyproject.toml

@@ -38,6 +38,7 @@ test = [
     "jsonschema",
     "pyyaml",
     "pytest-recording",
+    "deepdiff",
 ]
 dev = [
     "ruff==0.12.8",

src/anyvlm/anyvlm.py

@@ -0,0 +1,38 @@
+"""Client for cohort allele frequency (CAF) retrieval"""
+
+from anyvar.utils.types import VrsVariation
+from ga4gh.va_spec.base import CohortAlleleFrequencyStudyResult
+
+from anyvlm.storage.base_storage import Storage
+
+
+class AnyVLM:
+    """Client for cohort allele frequency (CAF) operations"""
+
+    def __init__(self, storage: Storage) -> None:
+        """Initialize AnyVLM client
+
+        :param storage: AnyVLM storage
+        """
+        self.storage = storage
+
+    def get_caf_for_variations(
+        self, vrs_variations: list[VrsVariation]
+    ) -> list[CohortAlleleFrequencyStudyResult]:
+        """Retrieve Cohort Allele Frequency data for given VRS Variations
+
+        :param vrs_variations: List of VRS Variations to get CAF data for
+        :return: list of CAFs for VRS Variations. Will use VRS Variation for focusAllele
+        """
+        vrs_variations_map: dict[str, VrsVariation] = {
+            vrs_variation.id: vrs_variation for vrs_variation in vrs_variations
+        }
+
+        cafs: list[CohortAlleleFrequencyStudyResult] = self.storage.get_caf_by_vrs_ids(
+            list(vrs_variations_map)
+        )
+
+        for caf in cafs:
+            caf.focusAllele = vrs_variations_map[caf.focusAllele.root]
+
+        return cafs

src/anyvlm/functions/get_caf.py

@@ -1,19 +1,29 @@
 """Perform search against variant(s) contained by an AnyVar node, and construct cohort allele frequency model(s)"""
 
+from anyvar.utils.types import VrsVariation
 from ga4gh.va_spec.base.core import CohortAlleleFrequencyStudyResult
 
 from anyvlm.anyvar.base_client import BaseAnyVarClient
+from anyvlm.anyvlm import AnyVLM
 
 
 def get_caf(
-    av: BaseAnyVarClient, accession_id: str, start: int, end: int
+    anyvar: BaseAnyVarClient,
+    anyvlm: AnyVLM,
+    accession_id: str,
+    start: int,
+    end: int,
 ) -> list[CohortAlleleFrequencyStudyResult]:
     """Retrieve Cohort Allele Frequency data for all known variants matching provided search params
 
-    :param av: AnyVar client
+    :param anyvar: AnyVar client (variant lookup)
+    :param anyvlm: AnyVLM client (caf lookup)
     :param accession_id: ID for sequence to search upon
     :param start: start of range search
     :param end: end of range to search
     :return: list of CAFs contained in search interval
     """
-    raise NotImplementedError
+    vrs_variations: list[VrsVariation] = anyvar.search_by_interval(
+        accession_id, start, end
+    )
+    return anyvlm.get_caf_for_variations(vrs_variations)

src/anyvlm/storage/base_storage.py

@@ -34,3 +34,14 @@ def add_allele_frequency(self, caf: CohortAlleleFrequencyStudyResult) -> None:
 
         :param caf: Cohort allele frequency study result object to insert into the DB
         """
+
+    @abstractmethod
+    def get_caf_by_vrs_ids(
+        self, vrs_ids: list[str]
+    ) -> list[CohortAlleleFrequencyStudyResult]:
+        """Retrieve cohort allele frequency study results by VRS IDs
+
+        :param vrs_ids: List of VRS variation IDs
+        :return: List of cohort allele frequency study results matching given VRS
+            variation IDs. Will use iriReference for focusAllele
+        """

src/anyvlm/storage/postgres.py

@@ -1,7 +1,7 @@
 """Provide PostgreSQL-based storage implementation."""
 
 from ga4gh.va_spec.base import CohortAlleleFrequencyStudyResult
-from sqlalchemy import create_engine, delete
+from sqlalchemy import create_engine, delete, select
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.orm import sessionmaker
 
@@ -46,3 +46,26 @@ def add_allele_frequency(self, caf: CohortAlleleFrequencyStudyResult) -> None:
 
         with self.session_factory() as session, session.begin():
             session.execute(stmt, db_entity.to_dict())
+
+    def get_caf_by_vrs_ids(
+        self, vrs_ids: list[str]
+    ) -> list[CohortAlleleFrequencyStudyResult]:
+        """Retrieve cohort allele frequency study results by VRS IDs
+
+        :param vrs_ids: List of VRS variation IDs
+        :return: List of cohort allele frequency study results matching given VRS
+            variation IDs. Will use iriReference for focusAllele
+        """
+        cafs: list[CohortAlleleFrequencyStudyResult] = []
+        with self.session_factory() as session:
+            stmt = (
+                select(orm.AlleleFrequencyData)
+                .where(orm.AlleleFrequencyData.vrs_id.in_(vrs_ids))
+                .limit(self.MAX_ROWS)
+            )
+            db_objects = session.scalars(stmt).all()
+
+            for db_object in db_objects:
+                caf = mapper_registry.from_db_entity(db_object)
+                cafs.append(caf)
+        return cafs

src/anyvlm/utils/__init__.py

@@ -0,0 +1 @@
+"""Provide utilities."""

src/anyvlm/utils/types.py

@@ -0,0 +1,18 @@
+"""Provide helpful type definitions, references, and type-based operations."""
+
+from pydantic import BaseModel
+
+
+class AncillaryResults(BaseModel):
+    """Define model for Ancillary Results"""
+
+    homozygotes: int
+    heterozygotes: int
+    hemizygotes: int
+    consequence: str
+
+
+class QualityMeasures(BaseModel):
+    """Define model for Quality Measures"""
+
+    qcFilters: list[str]  # noqa: N815

tests/conftest.py

@@ -3,12 +3,14 @@
 from pathlib import Path
 
 import pytest
+from anyvar.anyvar import create_storage, create_translator
 from dotenv import load_dotenv
 from ga4gh.core.models import iriReference
 from ga4gh.va_spec.base import CohortAlleleFrequencyStudyResult, StudyGroup
 from ga4gh.vrs import models
 from pydantic import BaseModel
 
+from anyvlm.anyvar.python_client import PythonAnyVarClient
 from anyvlm.storage import orm
 from anyvlm.storage.base_storage import Storage
 from anyvlm.storage.postgres import PostgresObjectStore
@@ -68,17 +70,42 @@ def vcr_config():
 
 
 @pytest.fixture(scope="session")
-def postgres_uri():
+def anyvlm_anyvar_postgres_uri():
+    return environ.get(
+        "ANYVLM_ANYVAR_TEST_STORAGE_URI",
+        "postgresql://postgres:postgres@localhost:5432/anyvlm_anyvar_test",
+    )
+
+
+@pytest.fixture
+def anyvar_python_client(anyvlm_anyvar_postgres_uri: str) -> PythonAnyVarClient:
+    storage = create_storage(anyvlm_anyvar_postgres_uri)
+    storage.wipe_db()
+    translator = create_translator()
+    return PythonAnyVarClient(translator, storage)
+
+
+@pytest.fixture
+def anyvar_populated_python_client(
+    anyvar_python_client: PythonAnyVarClient, alleles: dict
+):
+    for allele_fixture in alleles.values():
+        anyvar_python_client.put_objects([models.Allele(**allele_fixture["variation"])])
+    return anyvar_python_client
+
+
+@pytest.fixture(scope="session")
+def anyvlm_postgres_uri():
     return environ.get(
         "ANYVLM_TEST_STORAGE_URI",
         "postgresql://postgres:postgres@localhost:5432/anyvlm_test",
     )
 
 
 @pytest.fixture
-def postgres_storage(postgres_uri: str):
+def postgres_storage(anyvlm_postgres_uri: str):
     """Reset storage state after each test case"""
-    storage = PostgresObjectStore(postgres_uri)
+    storage = PostgresObjectStore(anyvlm_postgres_uri)
     yield storage
     storage.wipe_db()
 

tests/integration/functions/test_get_caf.py

@@ -0,0 +1,100 @@
+"""Test that get_caf function works correctly"""
+
+import pytest
+from deepdiff import DeepDiff
+from ga4gh.core.models import iriReference
+from ga4gh.va_spec.base import CohortAlleleFrequencyStudyResult
+
+from anyvlm.anyvar.python_client import PythonAnyVarClient
+from anyvlm.anyvlm import AnyVLM
+from anyvlm.functions.get_caf import get_caf
+from anyvlm.storage.postgres import PostgresObjectStore
+
+POS = 2781760
+REFGET_AC = "SQ.8_liLu1aycC0tPQPFmUaGXJLDs5SbPZ5"
+GA4GH_SEQ_ID = f"ga4gh:{REFGET_AC}"
+
+
+@pytest.fixture
+def alleles_to_add(alleles: dict):
+    """Create test fixture for alleles whose sequence reference matches REFGET_AC"""
+    return [
+        value["variation"]
+        for value in alleles.values()
+        if value["location"]["sequenceReference"]["refgetAccession"] == REFGET_AC
+    ]
+
+
+@pytest.fixture
+def alleles_in_range(alleles_to_add):
+    """Create test fixture for alleles overlapping POS"""
+    return [
+        variation
+        for variation in alleles_to_add
+        if variation["location"]["start"] <= POS and variation["location"]["end"] >= POS
+    ]
+
+
+@pytest.fixture
+def populated_postgres_storage(
+    postgres_storage: PostgresObjectStore,
+    alleles_to_add: list[dict],
+    caf_iri: CohortAlleleFrequencyStudyResult,
+):
+    for variation in alleles_to_add:
+        caf_copy = caf_iri.model_copy(deep=True)
+        caf_copy.focusAllele = iriReference(root=variation["id"])
+        postgres_storage.add_allele_frequency(caf_copy)
+    return postgres_storage
+
+
+@pytest.fixture
+def anyvlm_populated_client(populated_postgres_storage):
+    """Define test fixture for anyvlm"""
+    return AnyVLM(populated_postgres_storage)
+
+
+@pytest.fixture
+def expected_cafs(caf_iri, alleles_in_range):
+    cafs = []
+    for variation in alleles_in_range:
+        new_caf = caf_iri.model_copy(deep=True)
+        new_caf.focusAllele = variation
+        cafs.append(new_caf)
+    return cafs
+
+
+def test_get_caf_results_returned(
+    anyvar_populated_python_client: PythonAnyVarClient,
+    anyvlm_populated_client: AnyVLM,
+    expected_cafs: list[CohortAlleleFrequencyStudyResult],
+):
+    """Test get_caf when results are expected"""
+    cafs = get_caf(
+        anyvar_populated_python_client,
+        anyvlm_populated_client,
+        GA4GH_SEQ_ID,
+        POS,
+        POS,
+    )
+    diff = DeepDiff(
+        [caf.model_dump(exclude_none=True) for caf in cafs],
+        [caf.model_dump(exclude_none=True) for caf in expected_cafs],
+        ignore_order=True,
+    )
+    assert diff == {}
+
+
+def test_get_caf_no_results(
+    anyvar_populated_python_client: PythonAnyVarClient,
+    anyvlm_populated_client: AnyVLM,
+):
+    """Test get_caf when no results are expected"""
+    cafs = get_caf(
+        anyvar_populated_python_client,
+        anyvlm_populated_client,
+        "GRCh45.p1:Y",
+        POS,
+        POS,
+    )
+    assert cafs == []

tests/integration/storage/test_postgres_integration.py

@@ -6,10 +6,12 @@
 from anyvlm.storage.postgres import PostgresObjectStore
 
 
-def test_db_lifecycle(postgres_uri: str, caf_iri: CohortAlleleFrequencyStudyResult):
+def test_db_lifecycle(
+    anyvlm_postgres_uri: str, caf_iri: CohortAlleleFrequencyStudyResult
+):
     """Test that DB lifecycle works correctly"""
     # set up and populate DB
-    storage = PostgresObjectStore(postgres_uri)
+    storage = PostgresObjectStore(anyvlm_postgres_uri)
     caf_rows = return_cafs(storage)
     assert caf_rows == []