Merge pull request #2594 from cta-observatory/update-exception-handling

Update exception handling in tools

docs/changes/2594.optimization.rst

@@ -0,0 +1,4 @@
+Update exception handling in tools
+
+- Add a possibility to handle custom exception in ``Tool.run()``
+  with the preservation of the exit code.

src/ctapipe/core/tests/test_tool.py

@@ -381,13 +381,40 @@ 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."
+
+        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."
+
+        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"""

src/ctapipe/core/tests/test_traits.py

@@ -351,11 +351,11 @@ class MyTool(Tool):
         run_tool(tool, ["--MyTool.energy=5 m"], raises=True)
 
     captured = capsys.readouterr()
-    assert (
-        captured.err.split(":")[-1]
-        == f" Given quantity is of physical type {u.get_physical_type(5 * u.m)}."
-        + f" Expected {u.physical.energy}.\n"
+    expected = (
+        f" Given quantity is of physical type {u.get_physical_type(5 * u.m)}."
+        f" Expected {u.physical.energy}.\n"
     )
+    assert expected in captured.err
 
 
 def test_quantity_none():

src/ctapipe/core/tool.py

@@ -87,6 +87,9 @@ class Tool(Application):
     executing. This happens after the ``finish`` method has run or
     in case of exceptions.
 
+    User-defined code can raise custom exceptions both in the components
+    or in the tool methods. If these custom exceptions have an ``exit_code`` attribute,
+    it will be propagated to the final exit code of the tool.
 
     .. code:: python
 
@@ -403,6 +406,7 @@ def run(self, argv=None, raises=False):
         #  https://tldp.org/LDP/abs/html/exitcodes.html
 
         exit_status = 0
+        current_exception = None
 
         with self._exit_stack:
             try:
@@ -429,46 +433,47 @@ def run(self, argv=None, raises=False):
                 self.log.info("Finished: %s", self.name)
                 Provenance().finish_activity(activity_name=self.name)
             except (ToolConfigurationError, TraitError) as err:
+                current_exception = err
                 self.log.error("%s", err)
                 self.log.error("Use --help for more info")
                 exit_status = 2  # wrong cmd line parameter
                 Provenance().finish_activity(
                     activity_name=self.name, status="error", exit_code=exit_status
                 )
-                if raises:
-                    raise
             except KeyboardInterrupt:
                 self.log.warning("WAS INTERRUPTED BY CTRL-C")
                 exit_status = 130  # Script terminated by Control-C
                 Provenance().finish_activity(
                     activity_name=self.name, status="interrupted", exit_code=exit_status
                 )
             except Exception as err:
+                current_exception = err
+                exit_status = getattr(err, "exit_code", 1)
                 self.log.exception("Caught unexpected exception: %s", err)
-                exit_status = 1  # any other error
                 Provenance().finish_activity(
                     activity_name=self.name, status="error", exit_code=exit_status
                 )
-                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,
-                        )
+                if exit_status == 0:
+                    # Finish normally
+                    Provenance().finish_activity(activity_name=self.name)
+                else:
+                    # Finish with error
+                    current_exception = err
+                    self.log.critical(
+                        "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()
+                if raises and current_exception:
+                    raise current_exception
 
         self.exit(exit_status)