Repository: joerick/pyinstrument
Pull request #372: Fix %%pyinstrument magic to handle KeyboardInterrupt
Author: login: Darshan808
Created: 2025-05-03T06:17:56Z | Merged: 2025-05-23T15:38:39Z

## Fixes https://github.com/joerick/pyinstrument/issues/371

This PR ensures that the `%%pyinstrument` cell magic properly handles `KeyboardInterrupt`.

Previously, when a user interrupted execution (e.g., via Ctrl-C), the signal was not handled correctly, allowing subsequent cells to run unintentionally.

Now, the magic exits with `SIGINT` using `sys.exit(SIGINT)`, preventing further execution and aligning behavior with standard cell execution.

Logs on keyboard interrupt
```
---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[6], line 1
----> 1 sleep(10)

KeyboardInterrupt: 
An exception has occurred, use %tb to see the full traceback.

SystemExit: 2
```


Diff:
diff --git a/pyinstrument/magic/magic.py b/pyinstrument/magic/magic.py
index dc18306d..e9570f2d 100644
--- a/pyinstrument/magic/magic.py
+++ b/pyinstrument/magic/magic.py
@@ -40,6 +40,10 @@ def _get_active_profiler():
     return _active_profiler
 
 
+class InterruptSilently(Exception):
+    """Exception used to interrupt execution without showing traceback"""
+
+
 @magics_class
 class PyinstrumentMagic(Magics):
     def __init__(self, shell):
@@ -241,6 +245,24 @@ def pyinstrument(self, line, cell=None):
             )
             return
 
+        # If a KeyboardInterrupt occurred during the magic execution,
+        # raise an exception to prevent further executions.
+        if isinstance(cell_result.error_in_exec, KeyboardInterrupt):
+            # The traceback is already shown during the cell execution above, so we
+            # don't re-raise the exception directly.
+            old_custom_tb = ip.CustomTB
+            old_custom_exceptions = ip.custom_exceptions
+
+            def _silent_exception_handler(self, etype, value, tb, tb_offset=None):
+                # restore the original handlers
+                ip.CustomTB = old_custom_tb
+                ip.custom_exceptions = old_custom_exceptions
+                # swallow the InterruptSilently entirely
+
+            # install our silent handler
+            ip.set_custom_exc((InterruptSilently,), _silent_exception_handler)
+            raise InterruptSilently()
+
         html_config = compute_render_options(
             args, renderer_class=HTMLRenderer, unicode_support=True, color_support=True
         )
@@ -283,7 +305,7 @@ def run_cell_async(self, ip, code):
             asyncio.set_event_loop(loop)
             coro = ip.run_cell_async(code)
             future = asyncio.run_coroutine_threadsafe(coro, loop)
-            return future.result().result
+            return future.result()
         finally:
             loop.call_soon_threadsafe(loop.stop)
             asyncio.set_event_loop(old_loop)
diff --git a/test/test_ipython_magic.py b/test/test_ipython_magic.py
index e6ea5414..99887608 100644
--- a/test/test_ipython_magic.py
+++ b/test/test_ipython_magic.py
@@ -1,7 +1,13 @@
+import signal
+import textwrap
 from test.fake_time_util import fake_time
+from threading import Thread
+from time import sleep
 
 import pytest
 
+from pyinstrument.magic.magic import InterruptSilently
+
 # note: IPython should be imported within each test. Importing it in our tests
 # seems to cause problems with subsequent tests.
 
@@ -77,6 +83,40 @@ def test_magic_no_variable_expansion(ip, capsys):
     assert "hello 5" not in captured.out
 
 
+@pytest.mark.ipythonmagic
+def test_pyinstrument_handles_interrupt_silently(ip, capsys):
+    thread = Thread(target=_interrupt_after_1s)
+    thread.start()
+    # expect our custom exception to bubble up
+    with pytest.raises(InterruptSilently):
+        ip.run_cell_magic("pyinstrument", "", "from time import sleep; sleep(2)")
+
+    thread.join()
+    # nothing should have hit stderr
+    _, err = capsys.readouterr()
+    assert err.strip() == ""
+
+
+@pytest.mark.ipythonmagic
+def test_async_cell_with_pyinstrument(ip, capsys):
+    ip.run_cell_magic(
+        "pyinstrument",
+        line="--async_mode=enabled",
+        cell=textwrap.dedent(
+            """
+            import asyncio
+            async def function_a():
+                await asyncio.sleep(0.1)
+                return 42
+            a = await function_a()
+            print("a:", a)
+            """
+        ),
+    )
+    stdout, stderr = capsys.readouterr()
+    assert "a: 42" in stdout
+
+
 # Utils #
 
 
@@ -87,6 +127,11 @@ def session_ip():
     yield start_ipython()
 
 
+def _interrupt_after_1s():
+    sleep(1)
+    signal.raise_signal(signal.SIGINT)
+
+
 @pytest.fixture(scope="function")
 def ip(session_ip):
     session_ip.run_line_magic(magic_name="load_ext", line="pyinstrument")


Reviews:
user: login: joerick
body: Thanks @Darshan808 !
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-05T14:44:41Z
user: login: krassowski
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T09:58:43Z
user: login: krassowski
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T09:15:18Z
user: login: krassowski
body: This works for me, tested locally - thank you @Darshan808!

Also, thank you @joerick for looking into this!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T12:07:31Z
user: login: joerick
body: This looks great, thanks. I'll just try it out locally...

state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-23T13:34:43Z
user: login: joerick
body: Works for me!
state: APPROVED
author_association: OWNER
submitted_at: 2025-05-23T15:35:01Z

Review comments:
diff_hunk: @@ -241,6 +243,9 @@ def pyinstrument(self, line, cell=None):
             )
             return
 
+        if isinstance(cell_result.error_in_exec, KeyboardInterrupt):
+            sys.exit(SIGINT)
+
path: pyinstrument/magic/magic.py
user: login: joerick
body: Could you add a comment here, explaining why it's necessary?
created_at: 2025-05-05T14:44:34Z
author_association: OWNER
diff_hunk: @@ -241,6 +243,13 @@ def pyinstrument(self, line, cell=None):
             )
             return
 
+        # If a KeyboardInterrupt occurred during the magic execution,
+        # send a SIGINT signal to prevent further executions.
+        if isinstance(cell_result.error_in_exec, KeyboardInterrupt):
+            # The traceback is already shown during the cell execution above, so we
+            # don't re-raise the exception.
+            sys.exit(SIGINT)
path: pyinstrument/magic/magic.py
user: login: krassowski
body: I found a way to have no additional traceback nor warnings shown. This requires temporarily registering a custom exception handler and raising a custom exception:

```suggestion
            old_custom_tb = ip.CustomTB
            old_custom_exceptions = ip.custom_exceptions

            def _silent_exception_handler(self, etype, value, tb, tb_offset=None):
                """No-op handler for silent exception."""
                # restore the original exception handler
                ip.CustomTB = old_custom_tb
                ip.custom_exceptions = old_custom_exceptions

            ip.set_custom_exc((InterruptSilently, ), _silent_exception_handler)

            raise InterruptSilently()
```

where:

```python

class InterruptSilently(Exception):
    """Exception used to interrupt execution without showing traceback"""
```

This results in much cleaner:

![image](https://github.com/user-attachments/assets/7efaf4d9-c7a3-493c-ba8d-dfa4731c2bce)

created_at: 2025-05-09T09:58:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2,9 +2,11 @@
 
 import asyncio
 import html
+import sys
 import threading
 import urllib.parse
 from ast import parse
+from signal import SIGINT
path: pyinstrument/magic/magic.py
user: login: krassowski
body: No longer used:

```suggestion
import threading
import urllib.parse
from ast import parse
```
created_at: 2025-05-13T09:14:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -241,6 +247,24 @@ def pyinstrument(self, line, cell=None):
             )
             return
 
+        # If a KeyboardInterrupt occurred during the magic execution,
+        # send a SIGINT signal to prevent further executions.
path: pyinstrument/magic/magic.py
user: login: krassowski
body: ```suggestion
        # raise an exception to prevent further executions.
```
created_at: 2025-05-13T09:14:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,7 +1,13 @@
+import signal
+import sys
path: test/test_ipython_magic.py
user: login: krassowski
body: No longer used

```suggestion
```
created_at: 2025-05-13T09:15:08Z
author_association: CONTRIBUTOR

PR comments:
user: login: Darshan808
created_at: 2025-05-03T07:22:29Z
body: Both re-raising the `KeyboardInterrupt` and using `sys.exit(signal.SIGINT)` successfully prevent the execution of subsequent cells. However, `ip.run_cell()` already logs the error and traceback. Re-raising the exception results in duplicated and cluttered output with an additional traceback.

Using `sys.exit(signal.SIGINT)` provides a cleaner exit without redundant tracebacks.
author_association: CONTRIBUTOR
user: login: krassowski
created_at: 2025-05-09T09:49:17Z
body: @Darshan808 testing this locally I see:

```
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[4], line 1
----> 1 sleep(10)
      2 my_important_data = [
      3     # what did I do wrong here?
      4     "very",
      5     "important",
      6     "data"
      7 ]

KeyboardInterrupt: 
An exception has occurred, use %tb to see the full traceback.

SystemExit: 2
/home/krassowski/.pyenv/versions/3.12.2/envs/pyinstrument/lib/python3.12/site-packages/IPython/core/interactiveshell.py:3678: UserWarning: To exit: use 'exit', 'quit', or Ctrl-D.
  warn("To exit: use 'exit', 'quit', or Ctrl-D.", stacklevel=1)
```

The warning comes from [this line in IPython](https://github.com/ipython/ipython/blob/5a3bf345ed7c78e5fab84faa81285e5fd9dfa3da/IPython/core/interactiveshell.py#L3678).

I've tried to see if we can suppress this warning, but alas it is not possible.
author_association: CONTRIBUTOR
user: login: krassowski
created_at: 2025-05-21T11:09:38Z
body: @joerick anything we could do to help with merging/releasing this?
author_association: CONTRIBUTOR
user: login: joerick
created_at: 2025-05-23T13:38:20Z
body: hmm, i'm getting an error when just running a cell here-

<img width="1148" alt="image" src="https://github.com/user-attachments/assets/103b3b5d-01e1-4017-9795-c73dfae26550" />

author_association: OWNER
user: login: joerick
created_at: 2025-05-23T15:38:35Z
body: No idea what's going on with the windows failures. They've gotta be unrelated to this PR, so I'll merge as-is.
author_association: OWNER