From 6cfd59d673728038ee18698051473f56ccc8c7a1 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 1/9] 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()

From 7d8ead3401e2f7030c78729970f47f48c1c321d8 Mon Sep 17 00:00:00 2001
From: "mykhailo.dalchenko" <mykhailo.dalchenko@cta-consortium.org>
Date: Wed, 17 Jul 2024 15:34:35 +0200
Subject: [PATCH 2/9] Add changelog

---
 docs/changes/2594.optimization.rst | 5 +++++
 1 file changed, 5 insertions(+)
 create mode 100644 docs/changes/2594.optimization.rst

diff --git a/docs/changes/2594.optimization.rst b/docs/changes/2594.optimization.rst
new file mode 100644
index 00000000000..408857af226
--- /dev/null
+++ b/docs/changes/2594.optimization.rst
@@ -0,0 +1,5 @@
+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.

From 25fb2b9a4e8743414e1808e7bfff6402fdb44c14 Mon Sep 17 00:00:00 2001
From: "mykhailo.dalchenko" <mykhailo.dalchenko@cta-consortium.org>
Date: Wed, 17 Jul 2024 16:34:08 +0200
Subject: [PATCH 3/9] Add back SystemExit handling

---
 docs/changes/2594.optimization.rst      |  1 -
 src/ctapipe/core/tests/test_run_tool.py | 11 +++++++++++
 src/ctapipe/core/tool.py                | 16 ++++++++++++++++
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/docs/changes/2594.optimization.rst b/docs/changes/2594.optimization.rst
index 408857af226..0dd4bddce95 100644
--- a/docs/changes/2594.optimization.rst
+++ b/docs/changes/2594.optimization.rst
@@ -1,5 +1,4 @@
 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.
diff --git a/src/ctapipe/core/tests/test_run_tool.py b/src/ctapipe/core/tests/test_run_tool.py
index e42e97396a7..5b584818704 100644
--- a/src/ctapipe/core/tests/test_run_tool.py
+++ b/src/ctapipe/core/tests/test_run_tool.py
@@ -1,3 +1,4 @@
+import sys
 from subprocess import CalledProcessError
 
 import pytest
@@ -16,5 +17,15 @@ 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/tool.py b/src/ctapipe/core/tool.py
index e06f7e21816..aaeea0f28fe 100644
--- a/src/ctapipe/core/tool.py
+++ b/src/ctapipe/core/tool.py
@@ -470,6 +470,22 @@ def run(self, argv=None, raises=False):
                 )
                 if raises:
                     raise
+            except SystemExit as err:
+                exit_status = err.code
+                if (
+                    exit_status != 0
+                ):  # Do nothing if SystemExit was called with the exit code 0 (e.g. with -h option)
+                    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()

From 464e5a0e7434a315ef9186682018b33f18a6315e Mon Sep 17 00:00:00 2001
From: "mykhailo.dalchenko" <mykhailo.dalchenko@cta-consortium.org>
Date: Wed, 17 Jul 2024 17:08:34 +0200
Subject: [PATCH 4/9] Implement @maxnoe suggestions

- Simplify SystemExit handling
- Remove custom exception registration,
  but add a treatment of the custom exit codes when catching generic Exception
---
 docs/changes/2594.optimization.rst  |  2 +-
 src/ctapipe/core/tests/test_tool.py |  2 --
 src/ctapipe/core/tool.py            | 53 ++++++++++-------------------
 3 files changed, 19 insertions(+), 38 deletions(-)

diff --git a/docs/changes/2594.optimization.rst b/docs/changes/2594.optimization.rst
index 0dd4bddce95..c1045f4d33c 100644
--- a/docs/changes/2594.optimization.rst
+++ b/docs/changes/2594.optimization.rst
@@ -1,4 +1,4 @@
 Update exception handling in tools
 
-- Add a possibility to register custom exceptions and handle them in ``Tool.run()``
+- Add a possibility to handle custom exception in ``Tool.run()``
   with the preservation of the exit code.
diff --git a/src/ctapipe/core/tests/test_tool.py b/src/ctapipe/core/tests/test_tool.py
index 1306708bea5..a9322eb514b 100644
--- a/src/ctapipe/core/tests/test_tool.py
+++ b/src/ctapipe/core/tests/test_tool.py
@@ -390,7 +390,6 @@ class CustomErrorWithExitCode(Exception):
     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.")
@@ -398,7 +397,6 @@ def start(self):
     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.")
diff --git a/src/ctapipe/core/tool.py b/src/ctapipe/core/tool.py
index aaeea0f28fe..5e42e4ad4b5 100644
--- a/src/ctapipe/core/tool.py
+++ b/src/ctapipe/core/tool.py
@@ -102,10 +102,6 @@ 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)
@@ -185,8 +181,6 @@ 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"
@@ -449,22 +443,12 @@ 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
+                exit_status = getattr(err, "exit_code", 1)
+                if exit_status == 1:
+                    self.log.exception("Caught unexpected exception: %s", err)
+                else:
+                    self.log.error("Caught exception: %s", err)
                 Provenance().finish_activity(
                     activity_name=self.name, status="error", exit_code=exit_status
                 )
@@ -472,20 +456,19 @@ def run(self, argv=None, raises=False):
                     raise
             except SystemExit as err:
                 exit_status = err.code
-                if (
-                    exit_status != 0
-                ):  # Do nothing if SystemExit was called with the exit code 0 (e.g. with -h option)
-                    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
+                    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()

From d48fcb533e78559dcf4beff160b788ab3446bf2b Mon Sep 17 00:00:00 2001
From: "mykhailo.dalchenko" <mykhailo.dalchenko@cta-consortium.org>
Date: Wed, 17 Jul 2024 17:44:38 +0200
Subject: [PATCH 5/9] Add re-raise in SystemExit treatment for testing purposes

---
 src/ctapipe/core/tool.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/ctapipe/core/tool.py b/src/ctapipe/core/tool.py
index 5e42e4ad4b5..6325c909b53 100644
--- a/src/ctapipe/core/tool.py
+++ b/src/ctapipe/core/tool.py
@@ -460,6 +460,8 @@ def run(self, argv=None, raises=False):
                     # Finish normally
                     Provenance().finish_activity(activity_name=self.name)
                 else:
+                    if raises:
+                        raise
                     # Finish with error
                     self.log.critical(
                         "Caught SystemExit with exit code %s", exit_status

From 9da3fafcaa0f96a05494aa64148bc68b44bfd28e Mon Sep 17 00:00:00 2001
From: "mykhailo.dalchenko" <mykhailo.dalchenko@cta-consortium.org>
Date: Thu, 18 Jul 2024 16:44:55 +0200
Subject: [PATCH 6/9] Simplify re-raise logic

---
 src/ctapipe/core/tests/test_traits.py |  8 ++++----
 src/ctapipe/core/tool.py              | 12 ++++++------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/ctapipe/core/tests/test_traits.py b/src/ctapipe/core/tests/test_traits.py
index 06a7e4500e8..c84b0da185f 100644
--- a/src/ctapipe/core/tests/test_traits.py
+++ b/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():
diff --git a/src/ctapipe/core/tool.py b/src/ctapipe/core/tool.py
index 6325c909b53..40b055e70f6 100644
--- a/src/ctapipe/core/tool.py
+++ b/src/ctapipe/core/tool.py
@@ -403,6 +403,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,14 +430,13 @@ 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
@@ -444,6 +444,7 @@ def run(self, argv=None, raises=False):
                     activity_name=self.name, status="interrupted", exit_code=exit_status
                 )
             except Exception as err:
+                current_exception = err
                 exit_status = getattr(err, "exit_code", 1)
                 if exit_status == 1:
                     self.log.exception("Caught unexpected exception: %s", err)
@@ -452,17 +453,14 @@ def run(self, argv=None, raises=False):
                 Provenance().finish_activity(
                     activity_name=self.name, status="error", exit_code=exit_status
                 )
-                if raises:
-                    raise
             except SystemExit as err:
                 exit_status = err.code
                 if exit_status == 0:
                     # Finish normally
                     Provenance().finish_activity(activity_name=self.name)
                 else:
-                    if raises:
-                        raise
                     # Finish with error
+                    current_exception = err
                     self.log.critical(
                         "Caught SystemExit with exit code %s", exit_status
                     )
@@ -474,6 +472,8 @@ def run(self, argv=None, raises=False):
             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)
 

From 00fe47a16b2d75e21d35fc1cdd6d441e96829335 Mon Sep 17 00:00:00 2001
From: "mykhailo.dalchenko" <mykhailo.dalchenko@cta-consortium.org>
Date: Fri, 19 Jul 2024 14:35:35 +0200
Subject: [PATCH 7/9] Remove distinction between expected and unexpected
 exceptions in the log message

---
 src/ctapipe/core/tool.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/ctapipe/core/tool.py b/src/ctapipe/core/tool.py
index 40b055e70f6..c6aba8c463d 100644
--- a/src/ctapipe/core/tool.py
+++ b/src/ctapipe/core/tool.py
@@ -446,10 +446,7 @@ def run(self, argv=None, raises=False):
             except Exception as err:
                 current_exception = err
                 exit_status = getattr(err, "exit_code", 1)
-                if exit_status == 1:
-                    self.log.exception("Caught unexpected exception: %s", err)
-                else:
-                    self.log.error("Caught exception: %s", err)
+                self.log.exception("Caught unexpected exception: %s", err)
                 Provenance().finish_activity(
                     activity_name=self.name, status="error", exit_code=exit_status
                 )

From 023eacb69895cd97d77b4a6688ea0ac092d21fae Mon Sep 17 00:00:00 2001
From: "mykhailo.dalchenko" <mykhailo.dalchenko@cta-consortium.org>
Date: Tue, 30 Jul 2024 23:29:10 +0100
Subject: [PATCH 8/9] Update Tool documentation w.r.t Exceptions

---
 src/ctapipe/core/tool.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/ctapipe/core/tool.py b/src/ctapipe/core/tool.py
index c6aba8c463d..b8af0465dd4 100644
--- a/src/ctapipe/core/tool.py
+++ b/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
 

From 25a140082756bfb675d6b6b9b1c5868c1d7fb493 Mon Sep 17 00:00:00 2001
From: "mykhailo.dalchenko" <mykhailo.dalchenko@cta-consortium.org>
Date: Tue, 30 Jul 2024 23:49:11 +0100
Subject: [PATCH 9/9] Fix typo

---
 src/ctapipe/core/tool.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ctapipe/core/tool.py b/src/ctapipe/core/tool.py
index b8af0465dd4..cdc005a51c8 100644
--- a/src/ctapipe/core/tool.py
+++ b/src/ctapipe/core/tool.py
@@ -88,7 +88,7 @@ class Tool(Application):
     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,
+    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