From e11a75ccab2a79d412daa49229a1c2524fe57c5d Mon Sep 17 00:00:00 2001
From: John Vouvakis Manousakis <ioannis_vm@berkeley.edu>
Date: Fri, 18 Oct 2024 18:15:01 -0700
Subject: [PATCH] Add logging of uncaught exceptions to Logger class

- Modified Logger to capture uncaught exceptions and log traceback
  using `sys.excepthook`.
- Added `log_exception` method to handle uncaught exceptions.
- Updated test suite to verify exception logging using subprocess.
---
 pelicun/base.py                  | 38 +++++++++++++++++++++++++++
 pelicun/tests/basic/test_base.py | 45 ++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/pelicun/base.py b/pelicun/base.py
index 93d4ac46d..146026eed 100644
--- a/pelicun/base.py
+++ b/pelicun/base.py
@@ -47,6 +47,7 @@
 import json
 import pprint
 import sys
+import traceback
 import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional, TypeVar, overload
@@ -61,6 +62,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Callable
+    from types import TracebackType
 
     from pelicun.assessment import AssessmentBase
 
@@ -328,6 +330,10 @@ def __init__(
         self.reset_log_strings()
         control_warnings()
 
+        # Set sys.excepthook to handle uncaught exceptions
+        # https://docs.python.org/3/library/sys.html#sys.excepthook
+        sys.excepthook = self.log_exception
+
     def reset_log_strings(self) -> None:
         """Populate the string-related attributes of the logger."""
         if self.log_show_ms:
@@ -452,6 +458,38 @@ def print_system_info(self) -> None:
             prepend_timestamp=False,
         )
 
+    def log_exception(
+        self,
+        exc_type: type[BaseException],
+        exc_value: BaseException,
+        exc_traceback: TracebackType | None,
+    ) -> None:
+        """
+        Log uncaught exceptions and their traceback.
+
+        Parameters
+        ----------
+        exc_type : Type[BaseException]
+            The exception class.
+        exc_value : BaseException
+            The exception instance.
+        exc_traceback : Optional[TracebackType]
+            The traceback object representing the call stack at the point
+            where the exception occurred.
+
+        """
+        message = (
+            f"Unhandled exception occurred:\n"
+            f"{''.join(traceback.format_exception(exc_type, exc_value, exc_traceback))}"
+        )
+
+        if self.log_file is not None:
+            with Path(self.log_file).open('a', encoding='utf-8') as f:
+                f.write(message)
+
+        if self.print_log:
+            print(message, file=sys.stderr)  # noqa: T201
+
 
 # get the absolute path of the pelicun directory
 pelicun_path = Path(__file__).resolve().parent
diff --git a/pelicun/tests/basic/test_base.py b/pelicun/tests/basic/test_base.py
index 577df1430..01dd98bcf 100644
--- a/pelicun/tests/basic/test_base.py
+++ b/pelicun/tests/basic/test_base.py
@@ -44,6 +44,7 @@
 import argparse
 import io
 import re
+import subprocess  # noqa: S404
 import tempfile
 from contextlib import redirect_stdout
 from pathlib import Path
@@ -208,6 +209,50 @@ def test_logger_div() -> None:
             assert f.read()
 
 
+def test_logger_exception() -> None:
+    # Create a temporary directory for log files
+    temp_dir = tempfile.mkdtemp()
+
+    # Create a sample Python script that will raise an exception
+    test_script = Path(temp_dir) / 'test_script.py'
+    test_script_content = f"""
+import sys
+import traceback
+from pelicun.base import Logger
+
+log_file = "{temp_dir}/log.txt"
+
+log = Logger(log_file=log_file, verbose=True, log_show_ms=True, print_log=True)
+
+raise ValueError("Test exception in subprocess")
+"""
+
+    # Write the test script to the file
+    test_script.write_text(test_script_content)
+
+    # Use subprocess to run the script
+    process = subprocess.run(  # noqa: S603
+        ['python', str(test_script)],  # noqa: S607
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+
+    # Check that the process exited with an error
+    assert process.returncode == 1
+
+    # Check the stdout/stderr for the expected output
+    assert 'Test exception in subprocess' in process.stderr
+
+    # Check that the exception was logged in the log file
+    log_file = Path(temp_dir) / 'log.txt'
+    assert log_file.exists(), 'Log file was not created'
+    log_content = log_file.read_text()
+    assert 'Test exception in subprocess' in log_content
+    assert 'Traceback' in log_content
+    assert 'ValueError' in log_content
+
+
 def test_split_file_name() -> None:
     file_path = 'example.file.name.txt'
     name, extension = base.split_file_name(file_path)