From 2898f652b4e15e15340be76014d7495f6c97ce08 Mon Sep 17 00:00:00 2001
From: RalfG <ralfg@hotmail.be>
Date: Wed, 24 Jan 2024 14:46:39 +0100
Subject: [PATCH] Better handle exceptions when generating report

---
 ms2rescore/core.py         | 10 +++++++---
 ms2rescore/exceptions.py   |  6 ++++++
 ms2rescore/report/utils.py | 21 ++++++++++++++++-----
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/ms2rescore/core.py b/ms2rescore/core.py
index fffb1902..2a8429e6 100644
--- a/ms2rescore/core.py
+++ b/ms2rescore/core.py
@@ -11,6 +11,7 @@
 from ms2rescore.parse_spectra import get_missing_values
 from ms2rescore.report import generate
 from ms2rescore.rescoring_engines import mokapot, percolator
+from ms2rescore import exceptions
 
 logger = logging.getLogger(__name__)
 
@@ -148,9 +149,12 @@ def rescore(configuration: Dict, psm_list: Optional[PSMList] = None) -> None:
 
     # Write report
     if config["write_report"]:
-        generate.generate_report(
-            output_file_root, psm_list=psm_list, feature_names=feature_names, use_txt_log=True
-        )
+        try:
+            generate.generate_report(
+                output_file_root, psm_list=psm_list, feature_names=feature_names, use_txt_log=True
+            )
+        except exceptions.ReportGenerationError as e:
+            logger.error(e)
 
 
 def _write_feature_names(feature_names, output_file_root):
diff --git a/ms2rescore/exceptions.py b/ms2rescore/exceptions.py
index f8d1fa01..f2b06323 100644
--- a/ms2rescore/exceptions.py
+++ b/ms2rescore/exceptions.py
@@ -23,3 +23,9 @@ class ModificationParsingError(IDFileParsingError):
     """Identification file parsing error."""
 
     pass
+
+
+class ReportGenerationError(MS2RescoreError):
+    """Error while generating report."""
+
+    pass
diff --git a/ms2rescore/report/utils.py b/ms2rescore/report/utils.py
index 9ccc87ed..e8568e59 100644
--- a/ms2rescore/report/utils.py
+++ b/ms2rescore/report/utils.py
@@ -1,5 +1,6 @@
 """Utility functions for MS²Rescore report generation."""
 
+import logging
 from collections import defaultdict
 from csv import DictReader
 from pathlib import Path
@@ -9,6 +10,10 @@
 import psm_utils
 from mokapot import LinearConfidence, LinearPsmDataset, read_fasta
 
+from ms2rescore.exceptions import ReportGenerationError
+
+logger = logging.getLogger(__name__)
+
 
 def read_feature_names(feature_names_path: Path) -> dict:
     """Read feature names and mapping with feature generator from file."""
@@ -41,10 +46,11 @@ def get_confidence_estimates(
             "before_rescoring_score"
         ].astype(float)
     except KeyError as e:
-        raise KeyError(
+        raise ReportGenerationError(
             "No `before_rescoring_score` in PSM list provenance data. Ensure that the PSM list "
-            "was generated by MS²Rescore."
+            "was generated by MS²Rescore. Could not generate report."
         ) from e
+
     peptide = (
         pd.Series(psm_list["peptidoform"]).astype(str).str.replace(r"(/\d+$)", "", n=1, regex=True)
     )
@@ -55,8 +61,7 @@ def get_confidence_estimates(
             "before": score_before,
             "after": psm_list["score"],
         }
-    )
-    psms = psms.reset_index()
+    ).reset_index()
 
     if fasta_file:
         fasta = read_fasta(fasta_file)
@@ -73,6 +78,12 @@ def get_confidence_estimates(
         if fasta_file:
             lin_psm_dataset.add_proteins(fasta)
 
-        confidence[when] = lin_psm_dataset.assign_confidence()
+        try:
+            confidence[when] = lin_psm_dataset.assign_confidence()
+        except RuntimeError as e:
+            raise ReportGenerationError(
+                f"Error while assigning confidence estimates to PSMs ({when} rescoring). "
+                "Could not generate report."
+            ) from e
 
     return confidence["before"], confidence["after"]