From b51d5329ddae3b869189cd816b7d97f948e92784 Mon Sep 17 00:00:00 2001
From: John Vouvakis Manousakis <ioannis_vm@berkeley.edu>
Date: Thu, 31 Oct 2024 10:45:24 -0700
Subject: [PATCH] Handle multiple logger case.

Improve implementation by ensuring an exception is logged by all
instantiated loggers.

This is achieved by creating a module-level LoggerRegistry class that
keeps track of all instantiated loggers and defines a `log_exception`
class method which is assigned to `sys.excepthook`. The class method
documents the exception to all registered logger objects in the form
of a warning.
---
 pelicun/base.py                  | 46 +++++++++++++++++++++++++++++---
 pelicun/tests/basic/test_base.py | 38 +++++++++++++++++---------
 2 files changed, 68 insertions(+), 16 deletions(-)

diff --git a/pelicun/base.py b/pelicun/base.py
index 146026eed..ab442cadf 100644
--- a/pelicun/base.py
+++ b/pelicun/base.py
@@ -50,7 +50,7 @@
 import traceback
 import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, TypeVar, overload
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypeVar, overload
 
 import colorama
 import numpy as np
@@ -253,6 +253,44 @@ def rng(self) -> np.random.Generator:
         return self._rng
 
 
+# Define a module-level LoggerRegistry
+class LoggerRegistry:
+    """Registry to manage all logger instances."""
+
+    _loggers: ClassVar[list[Logger]] = []
+
+    # The @classmethod decorator allows this method to be called on
+    # the class itself, rather than on instances. It interacts with
+    # class-level data (like _loggers), enabling a single registry for
+    # all Logger instances without needing an object of LoggerRegistry
+    # itself.
+    @classmethod
+    def register(cls, logger: Logger) -> None:
+        """Register a logger instance."""
+        cls._loggers.append(logger)
+
+    @classmethod
+    def log_exception(
+        cls,
+        exc_type: type[BaseException],
+        exc_value: BaseException,
+        exc_traceback: TracebackType | None,
+    ) -> None:
+        """Log exceptions to all registered loggers."""
+        message = (
+            f"Unhandled exception occurred:"
+            f"\n"
+            f"{''.join(traceback.format_exception(exc_type, exc_value, exc_traceback))}"
+        )
+        for logger in cls._loggers:
+            logger.warning(message)
+
+
+# Update sys.excepthook to log exceptions in all loggers
+# https://docs.python.org/3/library/sys.html#sys.excepthook
+sys.excepthook = LoggerRegistry.log_exception
+
+
 class Logger:
     """Generate log files documenting execution events."""
 
@@ -330,9 +368,9 @@ 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
+        # Register the logger to the LoggerRegistry in order to
+        # capture raised exceptions.
+        LoggerRegistry.register(self)
 
     def reset_log_strings(self) -> None:
         """Populate the string-related attributes of the logger."""
diff --git a/pelicun/tests/basic/test_base.py b/pelicun/tests/basic/test_base.py
index 02922c13d..a80f06aee 100644
--- a/pelicun/tests/basic/test_base.py
+++ b/pelicun/tests/basic/test_base.py
@@ -221,16 +221,26 @@ def test_logger_exception() -> None:
     # 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 pathlib import Path
 from pelicun.base import Logger
 
-log_file = "{Path(temp_dir) / 'log.txt'}"
+log_file_A = Path("{temp_dir}") / 'log_A.txt'
+log_file_B = Path("{temp_dir}") / 'log_B.txt'
 
-log = Logger(log_file=log_file, verbose=True, log_show_ms=True, print_log=True)
+log_A = Logger(
+    log_file=log_file_A,
+    verbose=True,
+    log_show_ms=True,
+    print_log=True,
+)
+log_B = Logger(
+    log_file=log_file_B,
+    verbose=True,
+    log_show_ms=True,
+    print_log=True,
+)
 
-raise ValueError("Test exception in subprocess")
+raise ValueError('Test exception in subprocess')
 """
 
     # Write the test script to the file
@@ -248,15 +258,19 @@ def test_logger_exception() -> None:
     assert process.returncode == 1
 
     # Check the stdout/stderr for the expected output
-    assert 'Test exception in subprocess' in process.stderr
+    assert 'Test exception in subprocess' in process.stdout
 
     # 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
+    log_files = (
+        Path(temp_dir) / 'log_A_warnings.txt',
+        Path(temp_dir) / 'log_B_warnings.txt',
+    )
+    for log_file in log_files:
+        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: