New result class (#116)

* add new respose class

* return prediction

* minor

src/auto_acmg.py

@@ -1,21 +1,29 @@
 """Implementations of the PVS1 algorithm."""
 
-from typing import Optional
+from typing import List, Optional
 
 from loguru import logger
 
 from src.core.config import Config
-from src.criteria.auto_acmg import AutoACMGCriteria
-from src.defs.auto_acmg import ACMGCriteria
-from src.defs.auto_pvs1 import PVS1Prediction, PVS1PredictionPathMapping, PVS1PredictionSeqVarPath
+from src.criteria.auto_criteria import AutoACMGCriteria
+from src.defs.auto_acmg import ACMGPrediction, AutoACMGResult
+from src.defs.auto_pvs1 import PVS1Prediction, PVS1PredictionPathMapping, PVS1PredictionStrucVarPath
 from src.defs.exceptions import AutoAcmgBaseException, ParseError
 from src.defs.genome_builds import GenomeRelease
 from src.defs.seqvar import SeqVar, SeqVarResolver
 from src.defs.strucvar import StrucVar, StrucVarResolver
 from src.pvs1.auto_pvs1 import AutoPVS1
 
+#: Pathogenic PVS1 predictions of sequence variants.
+PVS1_POSITIVE_SEQVAR_PREDICTIONS = [
+    PVS1Prediction.PVS1,
+    PVS1Prediction.PVS1_Strong,
+    PVS1Prediction.PVS1_Moderate,
+    PVS1Prediction.PVS1_Supporting,
+]
+
 #: Criteria, which have no automatic prediction yet.
-NOT_IMPLEMENTED_CRITERIA = [
+NOT_IMPLEMENTED_CRITERIA: List[str] = [
     "PS2",  # De novo (both parents not tested)
     "PS3",  # Well-established in vitro or in vivo functional studies
     "PS4",  # Increased prevalence in affected individuals vs. controls
@@ -100,7 +108,7 @@ def resolve_variant(self) -> SeqVar | StrucVar | None:
             logger.error("An unexpected error occurred: {}", e)
             return None
 
-    def predict(self):
+    def predict(self) -> Optional[AutoACMGResult]:
         """
         Predict ACMG criteria for the specified variant.
 
@@ -119,7 +127,7 @@ def predict(self):
         variant = self.resolve_variant()
         if not variant:
             logger.error("Failed to resolve variant: {}", self.variant_name)
-            return
+            return None
 
         if isinstance(variant, SeqVar):
             logger.info(
@@ -128,30 +136,56 @@ def predict(self):
                 self.genome_release.name,
             )
             self.seqvar: SeqVar = variant
-            self.seqvar_pvs1_prediction: PVS1Prediction = PVS1Prediction.NotSet
-            self.seqvar_pvs1_prediction_path: PVS1PredictionSeqVarPath = (
-                PVS1PredictionSeqVarPath.NotSet
+            self.prediction: AutoACMGResult = AutoACMGResult()
+
+            # PP5 and BP6 criteria
+            self.prediction.pp5.prediction = ACMGPrediction.NotSet
+            self.prediction.pp5.comment = "PP5 prediction is deprecated."
+            self.prediction.bp6.prediction = ACMGPrediction.NotSet
+            self.prediction.bp6.comment = "BP6 prediction is deprecated."
+            logger.warning("Note, that PP5 and BP6 criteria are depricated and not predicted.")
+
+            # Not implemented criteria
+            for crit in NOT_IMPLEMENTED_CRITERIA:
+                setattr(
+                    getattr(self.prediction, crit.lower()),
+                    "prediction",
+                    ACMGPrediction.NotSet,
+                )
+                setattr(
+                    getattr(self.prediction, crit.lower()),
+                    "comment",
+                    f"{crit} prediction is not implemented.",
+                )
+            logger.warning(
+                "Some criteria are not implemented yet: {}",
+                NOT_IMPLEMENTED_CRITERIA,
             )
-            self.seqvar_criteria: ACMGCriteria = ACMGCriteria()
 
             # PVS1
             try:
                 logger.info("Predicting PVS1.")
                 pvs1 = AutoPVS1(self.seqvar, self.genome_release, config=self.config)
                 seqvar_prediction, seqvar_prediction_path = pvs1.predict()
                 if seqvar_prediction is None or seqvar_prediction_path is None:
-                    logger.error("Failed to predict PVS1 criteria.")
+                    raise AutoAcmgBaseException(
+                        "PVS1 prediction failed: prediction or prediction path is None."
+                    )
                 else:
-                    # Double check if the prediction path is indeed for sequence variant
-                    assert isinstance(seqvar_prediction_path, PVS1PredictionSeqVarPath)
-                    self.seqvar_pvs1_prediction = seqvar_prediction
-                    self.seqvar_pvs1_prediction_path = seqvar_prediction_path
-                    logger.info(
-                        "PVS1 prediction: {}.\nThe prediction path is:\n{}.",
-                        self.seqvar_pvs1_prediction.name,
-                        PVS1PredictionPathMapping[self.seqvar_pvs1_prediction_path],
+                    if seqvar_prediction == PVS1Prediction.NotSet:
+                        raise AutoAcmgBaseException("PVS1 prediction failed: prediction NotSet.")
+                    self.prediction.pvs1.prediction = (
+                        ACMGPrediction.Positive
+                        if seqvar_prediction in PVS1_POSITIVE_SEQVAR_PREDICTIONS
+                        else ACMGPrediction.Negative
+                    )
+                    self.prediction.pvs1.comment = (
+                        f"PVS1 strength: {seqvar_prediction.name}. "
+                        f"PVS1 prediction path: {PVS1PredictionPathMapping[seqvar_prediction_path]}."
                     )
             except AutoAcmgBaseException as e:
+                self.prediction.pvs1.prediction = ACMGPrediction.NotSet
+                self.prediction.pvs1.comment = "PVS1 prediction failed."
                 logger.error("Failed to predict PVS1 criteria. Error: {}", e)
 
             # Other criteria
@@ -162,52 +196,48 @@ def predict(self):
                 )
                 criteria_preciction = auto_criteria.predict()
                 if criteria_preciction is None:
-                    logger.error("Failed to predict other ACMG criteria.")
+                    raise AutoAcmgBaseException("Other ACMG criteria prediction failed.")
                 else:
-                    self.seqvar_criteria = criteria_preciction
-                    logger.info("Other ACMG criteria prediction:\n{}", self.seqvar_criteria)
-                    logger.warning(
-                        "Note, that PP5 and BP6 criteria are depricated and not predicted."
-                    )
-                    logger.warning(
-                        "Some criteria are not implemented yet. They are: {}",
-                        NOT_IMPLEMENTED_CRITERIA,
-                    )
+                    for criteria, prediction in criteria_preciction.model_dump().items():
+                        setattr(
+                            getattr(self.prediction, criteria.lower()),
+                            "prediction",
+                            ACMGPrediction.Positive if prediction else ACMGPrediction.Negative,
+                        )
             except AutoAcmgBaseException as e:
                 logger.error("Failed to predict other ACMG criteria. Error: {}", e)
 
+            logger.info("ACMG criteria prediction completed.")
+            return self.prediction
+
         elif isinstance(variant, StrucVar):
-            logger.info("Classification of structural variants is not implemented yet.")
-            # logger.info("Currently only PVS1 prediction is implemented for structural variants!")
+            logger.warning("Currently there's no prediction for structural variants!")
+            # logger.warning("Currently only PVS1 prediction is implemented for structural variants!")
             # logger.info(
             #     "Classifying ACMG criteria for structural variant {}, genome release: {}.",
             #     variant.user_repr,
             #     self.genome_release.name,
             # )
+            # self.strucvar: StrucVar = variant
+            # self.strucvar_prediction: PVS1Prediction = PVS1Prediction.NotSet  # type: ignore
+            # self.strucvar_prediction_path: PVS1PredictionStrucVarPath = PVS1PredictionStrucVarPath.NotSet  # type: ignore
+
             # # PVS1
             # try:
             #     logger.info("Predicting PVS1.")
-            #     self.strucvar: StrucVar = variant
-            #     self.strucvar_prediction: PVS1Prediction = PVS1Prediction.NotSet  # type: ignore
-            #     self.strucvar_prediction_path: PVS1PredictionStrucVarPath = PVS1PredictionStrucVarPath.NotSet  # type: ignore
-
             #     pvs1 = AutoPVS1(self.strucvar, self.genome_release)
             #     strucvar_prediction, strucvar_prediction_path = pvs1.predict()
             #     if strucvar_prediction is None or strucvar_prediction_path is None:
-            #         logger.error("Failed to predict PVS1 criteria.")
-            #         return
+            #         raise AutoAcmgBaseException(
+            #             "PVS1 prediction failed: prediction or prediction path is None."
+            #         )
             #     else:
-            #         # Double check if the prediction path is indeed for structural variant
-            #         assert isinstance(strucvar_prediction_path, PVS1PredictionStrucVarPath)
+            #         if strucvar_prediction == PVS1Prediction.NotSet:
+            #             raise AutoAcmgBaseException("PVS1 prediction failed: prediction NotSet.")
             #         self.strucvar_prediction = strucvar_prediction
             #         self.strucvar_prediction_path = strucvar_prediction_path
-            #         logger.info(
-            #             "PVS1 prediction for {}: {}.\n" "The prediction path is:\n{}.",
-            #             self.strucvar.user_repr,
-            #             self.strucvar_prediction.name,
-            #             PVS1PredictionPathMapping[self.strucvar_prediction_path],
-            #         )
             # except AutoAcmgBaseException as e:
             #     logger.error("Failed to predict PVS1 criteria. Error: {}", e)
             #     return
-        logger.info("ACMG criteria prediction completed.")
+
+        return None

src/defs/auto_acmg.py

@@ -1,5 +1,5 @@
-from enum import Enum, EnumMeta
-from typing import Any
+from enum import Enum, EnumMeta, auto
+from typing import Any, List, Optional
 
 from pydantic import BaseModel
 
@@ -129,3 +129,52 @@ class PP3BP4(BaseModel):
 
     PP3: bool = False
     BP4: bool = False
+
+
+class ACMGPrediction(AutoAcmgBaseEnum):
+    """ACMG prediction enumeration."""
+
+    NotSet = auto()
+    Positive = auto()
+    Negative = auto()
+
+
+class CriteriaPrediction(BaseModel):
+    """Criteria prediction."""
+
+    name: str
+    prediction: ACMGPrediction = ACMGPrediction.NotSet
+    comment: str = ""
+
+
+class AutoACMGResult(BaseModel):
+    """Response of the ACMG criteria prediction."""
+
+    pvs1: CriteriaPrediction = CriteriaPrediction(name="PVS1")
+    ps1: CriteriaPrediction = CriteriaPrediction(name="PS1")
+    ps2: CriteriaPrediction = CriteriaPrediction(name="PS2")
+    ps3: CriteriaPrediction = CriteriaPrediction(name="PS3")
+    ps4: CriteriaPrediction = CriteriaPrediction(name="PS4")
+    pm1: CriteriaPrediction = CriteriaPrediction(name="PM1")
+    pm2: CriteriaPrediction = CriteriaPrediction(name="PM2")
+    pm3: CriteriaPrediction = CriteriaPrediction(name="PM3")
+    pm4: CriteriaPrediction = CriteriaPrediction(name="PM4")
+    pm5: CriteriaPrediction = CriteriaPrediction(name="PM5")
+    pm6: CriteriaPrediction = CriteriaPrediction(name="PM6")
+    pp1: CriteriaPrediction = CriteriaPrediction(name="PP1")
+    pp2: CriteriaPrediction = CriteriaPrediction(name="PP2")
+    pp3: CriteriaPrediction = CriteriaPrediction(name="PP3")
+    pp4: CriteriaPrediction = CriteriaPrediction(name="PP4")
+    pp5: CriteriaPrediction = CriteriaPrediction(name="PP5")
+    ba1: CriteriaPrediction = CriteriaPrediction(name="BA1")
+    bs1: CriteriaPrediction = CriteriaPrediction(name="BS1")
+    bs2: CriteriaPrediction = CriteriaPrediction(name="BS2")
+    bs3: CriteriaPrediction = CriteriaPrediction(name="BS3")
+    bs4: CriteriaPrediction = CriteriaPrediction(name="BS4")
+    bp1: CriteriaPrediction = CriteriaPrediction(name="BP1")
+    bp2: CriteriaPrediction = CriteriaPrediction(name="BP2")
+    bp3: CriteriaPrediction = CriteriaPrediction(name="BP3")
+    bp4: CriteriaPrediction = CriteriaPrediction(name="BP4")
+    bp5: CriteriaPrediction = CriteriaPrediction(name="BP5")
+    bp6: CriteriaPrediction = CriteriaPrediction(name="BP6")
+    bp7: CriteriaPrediction = CriteriaPrediction(name="BP7")

tests/e2e/test_e2e_other_criteria.py

@@ -4,15 +4,16 @@
 
 import pytest
 
-from src.auto_acmg import AutoACMG, AutoACMGCriteria
+from src.auto_acmg import AutoACMG
 from src.core.config import Config
-from src.defs.auto_acmg import ACMGCriteria
+from src.criteria.auto_criteria import AutoACMGCriteria
+from src.defs.auto_acmg import ACMGCriteria, AutoACMGResult
 from src.defs.genome_builds import GenomeRelease
 from src.defs.seqvar import SeqVar
 from tests.utils import load_test_data
 
 #: Type for ACMG criteria test data.
-AcmgTestData = List[Tuple[str, GenomeRelease, ACMGCriteria, str]]
+AcmgTestData = List[Tuple[str, GenomeRelease, AutoACMGResult, str]]
 #: Test data.
 ACMG_TEST_DATA: AcmgTestData = load_test_data("tests/assets/e2e_variants/other_criteria.csv")
 

tests/test_auto_acmg.py

@@ -5,9 +5,9 @@
 
 from src.auto_acmg import AutoACMG
 from src.core.config import Config
-from src.criteria.auto_acmg import AutoACMGCriteria
+from src.criteria.auto_criteria import AutoACMGCriteria
 from src.criteria.auto_ps1_pm5 import AutoPS1PM5
-from src.defs.auto_acmg import PS1PM5, ACMGCriteria
+from src.defs.auto_acmg import PS1PM5, ACMGCriteria, AutoACMGResult
 from src.defs.exceptions import AutoAcmgBaseException, ParseError
 from src.defs.genome_builds import GenomeRelease
 from src.defs.seqvar import SeqVar, SeqVarResolver
@@ -147,7 +147,7 @@ def test_auto_acmg_predict_failure_pvs1(
     assert mock_seqvar_resolver.resolve_seqvar.called
     assert mock_auto_pvs1_failure.predict.called
     assert auto_acmg.seqvar == mock_seqvar
-    assert auto_acmg.seqvar_pvs1_prediction_path == PVS1PredictionSeqVarPath.NotSet
+    assert auto_acmg.prediction is not None
 
 
 def test_auto_acmg_predict_failure_criteria(
@@ -164,3 +164,5 @@ def test_auto_acmg_predict_failure_criteria(
     assert mock_seqvar_resolver.resolve_seqvar.called
     assert mock_auto_pvs1_success.predict.called
     assert mock_auto_criteria_failure.predict.called
+    assert auto_acmg.seqvar == mock_seqvar
+    assert auto_acmg.prediction is not None