From 5dd363d7c598a72e9fd22c358c0391acecb115f1 Mon Sep 17 00:00:00 2001
From: "mykhailo.dalchenko" <mykhailo.dalchenko@cta-consortium.org>
Date: Wed, 17 Jul 2024 15:21:56 +0200
Subject: [PATCH] Update exception handling in tools

- Remove SystemExit handling in `Tool.run()`
- Add a possibility to register custom exceptions
  and handle them in `Tool.run()` with the preservation of the exit code.
---
 src/ctapipe/core/tests/test_run_tool.py | 11 --------
 src/ctapipe/core/tests/test_tool.py     | 29 +++++++++++++++++++++
 src/ctapipe/core/tool.py                | 34 ++++++++++++++-----------
 3 files changed, 48 insertions(+), 26 deletions(-)

diff --git a/src/ctapipe/core/tests/test_run_tool.py b/src/ctapipe/core/tests/test_run_tool.py
index 5b584818704..e42e97396a7 100644
--- a/src/ctapipe/core/tests/test_run_tool.py
+++ b/src/ctapipe/core/tests/test_run_tool.py
@@ -1,4 +1,3 @@
-import sys
 from subprocess import CalledProcessError
 
 import pytest
@@ -17,15 +16,5 @@ def start(self):
     ret = run_tool(ErrorTool(), ["--non-existing-alias"], raises=False)
     assert ret == 2
 
-    class SysExitTool(Tool):
-        def setup(self):
-            pass
-
-        def start(self):
-            sys.exit(4)
-
-    ret = run_tool(SysExitTool(), raises=False)
-    assert ret == 4
-
     with pytest.raises(CalledProcessError):
         run_tool(ErrorTool(), ["--non-existing-alias"], raises=True)
diff --git a/src/ctapipe/core/tests/test_tool.py b/src/ctapipe/core/tests/test_tool.py
index 53dc5a18e9a..1306708bea5 100644
--- a/src/ctapipe/core/tests/test_tool.py
+++ b/src/ctapipe/core/tests/test_tool.py
@@ -381,13 +381,42 @@ class ToolBad(Tool):
         def start(self):
             raise ValueError("1 does not equal 0.")
 
+    class CustomErrorNoExitCode(Exception):
+        pass
+
+    class CustomErrorWithExitCode(Exception):
+        exit_code = 42
+
+    class ToolCustomExceptionNoExitCode(Tool):
+        name = "CustomException"
+        description = "This tool raises a custom exception without an exit code."
+        custom_exceptions = (CustomErrorNoExitCode,)
+
+        def start(self):
+            raise CustomErrorNoExitCode("This is a custom exception.")
+
+    class ToolCustomExceptionWithExitCode(Tool):
+        name = "CustomException"
+        description = "This tool raises a custom exception with a custom exit code."
+        custom_exceptions = (CustomErrorWithExitCode,)
+
+        def start(self):
+            raise CustomErrorWithExitCode("This is a custom exception.")
+
     assert run_tool(ToolGood(), raises=True) == 0
 
     assert run_tool(ToolBad(), raises=False) == 1
 
+    assert run_tool(ToolCustomExceptionNoExitCode(), raises=False) == 1
+
+    assert run_tool(ToolCustomExceptionWithExitCode(), raises=False) == 42
+
     with pytest.raises(ValueError):
         run_tool(ToolBad(), raises=True)
 
+    with pytest.raises(CustomErrorNoExitCode):
+        run_tool(ToolCustomExceptionNoExitCode(), raises=True)
+
 
 def test_exit_stack():
     """Test that components that are context managers are properly handled"""
diff --git a/src/ctapipe/core/tool.py b/src/ctapipe/core/tool.py
index e1f7c2e2746..e06f7e21816 100644
--- a/src/ctapipe/core/tool.py
+++ b/src/ctapipe/core/tool.py
@@ -102,6 +102,10 @@ class MyTool(Tool):
             # Which classes are registered for configuration
             classes = [MyComponent, AdvancedComponent, SecondaryMyComponent]
 
+            # Tuple with the custom exceptions the tool can raise (optional)
+            # If the exception has an exit_code attribute, it will be used
+            custom_exceptions = (MyCustomError, MyOtherError)
+
             # local configuration parameters
             iterations = Integer(5,help="Number of times to run",
                                  allow_none=False).tag(config=True)
@@ -181,6 +185,8 @@ def main():
 
     provenance_log = Path(directory_ok=False).tag(config=True)
 
+    custom_exceptions = ()
+
     @default("provenance_log")
     def _default_provenance_log(self):
         return self.name + ".provenance.log"
@@ -443,6 +449,19 @@ def run(self, argv=None, raises=False):
                 Provenance().finish_activity(
                     activity_name=self.name, status="interrupted", exit_code=exit_status
                 )
+            except self.custom_exceptions as err:
+                self.log.exception("Caught custom exception: %s", err)
+                if hasattr(err, "exit_code"):
+                    exit_status = err.exit_code
+                else:
+                    exit_status = 1
+                Provenance().finish_activity(
+                    activity_name=self.name,
+                    status="error",
+                    exit_code=exit_status,
+                )
+                if raises:
+                    raise
             except Exception as err:
                 self.log.exception("Caught unexpected exception: %s", err)
                 exit_status = 1  # any other error
@@ -451,21 +470,6 @@ def run(self, argv=None, raises=False):
                 )
                 if raises:
                     raise
-            except SystemExit as err:
-                exit_status = err.code
-                # Do nothing if SystemExit was called with the exit code 0 (e.g. with -h option)
-                if exit_status != 0:
-                    if raises:
-                        raise  # do not re-intercept in tests
-                    else:
-                        self.log.exception(
-                            "Caught SystemExit with exit code %s", exit_status
-                        )
-                        Provenance().finish_activity(
-                            activity_name=self.name,
-                            status="error",
-                            exit_code=exit_status,
-                        )
             finally:
                 if not {"-h", "--help", "--help-all"}.intersection(self.argv):
                     self.write_provenance()