add validation for ResultSet ids + add extra info to TODOs

Author: jennifer-bowser

src/anyvlm/functions/build_vlm_response.py

@@ -1,5 +1,6 @@
 """Craft a VlmResponse object from a list of CohortAlleleFrequencyStudyResults"""
 
+from ga4gh.core.models import MappableConcept
 from ga4gh.va_spec.base.core import CohortAlleleFrequencyStudyResult
 
 from anyvlm.schemas.vlm import (
@@ -9,6 +10,41 @@
     ResultSet,
     VlmResponse,
 )
+from anyvlm.utils.types import Zygosity
+
+
+def _extract_zygosity(caf_study_result: CohortAlleleFrequencyStudyResult) -> Zygosity:
+    """Extracts a zygosity from a CohortAlleleFrequencyStudyResult.
+    #TODO: This is my best guess at how we'll want to represent zygosity in our CAF results. We may need
+    to update this function during or after Issue #16 based on how we choose to implement our CAF data pull.
+
+    :param caf_study_result: The CohortAlleleFrequencyStudyResult whose zygosity we wish to determine.
+    :return: The `Zygosity` of the study result.
+    """
+    cohort_characteristics: list[MappableConcept] | None = (
+        caf_study_result.cohort.characteristics
+    )
+    if not (cohort_characteristics):
+        error_message: str = "Each CohortAlleleFrequencyStudyResult's 'cohort' field must have 'characteristics' set"
+        raise ValueError(error_message)
+
+    zygosity: str | None = next(
+        (
+            concept.name
+            for concept in cohort_characteristics
+            if getattr(concept, "conceptType", None) == "Zygosity"
+        ),
+        None,
+    )
+    if not zygosity:
+        error_message: str = "'CohortAlleleFrequencyStudyResult.cohort.characteristics' must contain a 'Zygosity' entry"
+        raise ValueError(error_message)
+
+    try:
+        return Zygosity(zygosity)
+    except ValueError as e:
+        error_message: str = f"Invalid zygosity provided: {zygosity}"
+        raise ValueError(error_message) from e
 
 
 def build_vlm_response(caf_data: list[CohortAlleleFrequencyStudyResult]) -> VlmResponse:
@@ -21,15 +57,17 @@ def build_vlm_response(caf_data: list[CohortAlleleFrequencyStudyResult]) -> VlmR
     if caf_data:
         total: int = sum(
             [caf_study_result.focusAlleleCount for caf_study_result in caf_data]
-        )  # TODO: I'm not sure this is the right field?
+        )  # TODO: I'm not sure this is the right field? Will need to verify during/after Issue #16.
         response_summary = ResponseSummary(exists=True, total=total)
 
         for caf_study_result in caf_data:
             result_sets.extend(
                 [
                     ResultSet(
                         exists=True,
-                        id=f"{HandoverType.id} {caf_study_result.cohort}",  # TODO - not sure that caf_study_result.cohort is the right field
+                        # TODO - HandoverType.id represents the ID of the node from which the dataset was pulled.
+                        # In the future, this ID should be set dynamically.
+                        id=f"{HandoverType.id} {_extract_zygosity(caf_study_result)}",
                         resultsCount=caf_study_result.focusAlleleCount,
                     )
                 ]

src/anyvlm/schemas/vlm.py

@@ -1,8 +1,11 @@
 """Schemas relating to VLM API."""
 
-from pydantic import BaseModel, Field
+from typing import Self
+
+from pydantic import BaseModel, Field, model_validator
 
 from anyvlm import __version__
+from anyvlm.utils.types import Zygosity
 
 # ruff: noqa: N815 (allows camelCase vars instead of snake_case)
 
@@ -88,3 +91,26 @@ class VlmResponse(BaseModel):
     meta: Meta = Meta()
     responseSummary: ResponseSummary
     response: ResponseField
+
+    @model_validator(mode="after")
+    def validate_resultset_ids(self) -> Self:
+        """Ensure each ResultSet.id is correctly constructed."""
+        handover_ids: list[str] = [
+            beaconHandover.handoverType.id for beaconHandover in self.beaconHandovers
+        ]
+
+        for result_set in self.response.resultSets:
+            node_id, zygosity = result_set.id.split(" ")
+
+            if node_id not in handover_ids:
+                error_message = f"Invalid ResultSet id - ids must be in form '<node_id> <zygosity>', but provided node_id of {node_id} does not match any `handoverType.id` provided in `self.beaconHandovers`"
+                raise ValueError(error_message)
+
+            try:
+                Zygosity(zygosity)
+            except ValueError as e:
+                valid_zygosity_values = {zygosity.value for zygosity in Zygosity}
+                error_message = f"Invalid ResultSet id - ids must be in form '<node_id> <zygosity>', but provided zygosity of {zygosity} is not found in allowable value set of: {', '.join(valid_zygosity_values)}"
+                raise ValueError(error_message) from e
+
+        return self

src/anyvlm/utils/types.py

@@ -74,3 +74,12 @@ def _normalize_chromosome_name(chromosome_name: str) -> str:
 
 
 ChromosomeName = Annotated[str, BeforeValidator(_normalize_chromosome_name)]
+
+
+class Zygosity(StrEnum):
+    """Allowable zygosity values as defined by the VLM protocol"""
+
+    HOMOZYGOUS = "Homozygous"
+    HETEROZYGOUS = "Heterozygous"
+    HEMIZYGOUS = "Hemizygous"
+    UNKNOWN = "Unknown Zygosity"