Merge pull request #319 from frantic/rethrow-inner-exceptions

Propagate `SystemExit` and `KeyboardInterrupt` from CLI

pyinstrument/__main__.py

@@ -298,6 +298,8 @@ def store_and_consume_remaining(
         f = sys.stdout
         should_close_f_after_writing = False
 
+    inner_exception = None
+
     # create the renderer
 
     try:
@@ -363,8 +365,8 @@ def store_and_consume_remaining(
         try:
             sys.argv[:] = argv
             exec(code, globs, None)
-        except (SystemExit, KeyboardInterrupt):
-            pass
+        except (SystemExit, KeyboardInterrupt) as e:
+            inner_exception = e
         finally:
             sys.argv[:] = old_argv
 
@@ -385,6 +387,11 @@ def store_and_consume_remaining(
         print("    pyinstrument --load-prev %s [options]" % report_identifier)
         print("")
 
+    if inner_exception:
+        # If the script raised an exception, re-raise it now to resume
+        # the normal Python exception handling (printing the traceback, etc.)
+        raise inner_exception
+
 
 def compute_render_options(
     options: CommandLineOptions, renderer_class: type[renderers.Renderer], output_file: TextIO

test/test_cmdline.py

@@ -60,7 +60,7 @@ def test_single_file_module_running(self, pyinstrument_invocation, tmp_path: Pat
 
     def test_running_yourself_as_module(self, pyinstrument_invocation):
         subprocess.check_call(
-            [*pyinstrument_invocation, "-m", "pyinstrument"],
+            [*pyinstrument_invocation, "-m", "pyinstrument", "--help"],
         )
 
     def test_path(self, pyinstrument_invocation, tmp_path: Path, monkeypatch):
@@ -295,3 +295,16 @@ def test_binary_output(self, pyinstrument_invocation, tmp_path: Path):
 
         stats = pstats.Stats(str(output_file))
         assert stats
+
+    def test_program_exit_code(self, pyinstrument_invocation, tmp_path: Path):
+        exit_1_py = tmp_path / "exit_1.py"
+        exit_1_py.write_text("""import sys; sys.exit(1)""")
+
+        retcode = subprocess.call(
+            [
+                *pyinstrument_invocation,
+                str(exit_1_py),
+            ],
+        )
+
+        assert retcode == 1