fix ctrl+c handling, #37

Author: egorsmkv

python/rsloop/_run.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import asyncio as __asyncio
+import signal as __signal
 import sys as __sys
+import threading as __threading
 import typing as __typing
 
 from ._loop_compat import Loop
@@ -11,6 +13,10 @@
 _PREVIOUS_EVENT_LOOP_POLICY: __typing.Optional[__asyncio.AbstractEventLoopPolicy] = None
 
 
+def _noop() -> None:
+    pass
+
+
 def new_event_loop() -> Loop:
     return Loop()
 
@@ -48,6 +54,25 @@ def uninstall() -> None:
     _PREVIOUS_EVENT_LOOP_POLICY = None
 
 
+class _SigintHandler:
+    def __init__(self, loop: Loop, main_task: __asyncio.Task[__typing.Any]) -> None:
+        self._loop = loop
+        self._main_task = main_task
+        self.interrupt_count = 0
+
+    def __call__(
+        self,
+        signum: int,
+        frame: __typing.Optional[__typing.Any],
+    ) -> None:
+        self.interrupt_count += 1
+        if self.interrupt_count == 1 and not self._main_task.done():
+            self._main_task.cancel()
+            self._loop.call_soon_threadsafe(_noop)
+            return
+        raise KeyboardInterrupt()
+
+
 if __typing.TYPE_CHECKING:
 
     def run(
@@ -86,7 +111,31 @@ async def wrapper():
             __asyncio.set_event_loop(loop)
             if debug is not None:
                 loop.set_debug(debug)
-            return loop.run_until_complete(wrapper())
+            main_task = loop.create_task(wrapper())
+            sigint_handler = None
+            if (
+                __threading.current_thread() is __threading.main_thread()
+                and __signal.getsignal(__signal.SIGINT) is __signal.default_int_handler
+            ):
+                sigint_handler = _SigintHandler(loop, main_task)
+                try:
+                    __signal.signal(__signal.SIGINT, sigint_handler)
+                except ValueError:
+                    sigint_handler = None
+            try:
+                return loop.run_until_complete(main_task)
+            except __asyncio.CancelledError:
+                if sigint_handler is not None and sigint_handler.interrupt_count > 0:
+                    uncancel = getattr(main_task, "uncancel", None)
+                    if uncancel is None or uncancel() == 0:
+                        raise KeyboardInterrupt()
+                raise
+            finally:
+                if (
+                    sigint_handler is not None
+                    and __signal.getsignal(__signal.SIGINT) is sigint_handler
+                ):
+                    __signal.signal(__signal.SIGINT, __signal.default_int_handler)
         finally:
             try:
                 __cancel_all_tasks(loop)

tests/test_run.py

@@ -102,6 +102,47 @@ async def main() -> str:
 
         self.assertEqual(rsloop.run(main()), "ok")
 
+    def test_run_fallback_handles_sigint(self) -> None:
+        script = r"""
+import asyncio
+import signal
+import sys
+import threading
+
+import rsloop
+import rsloop._run as rsloop_run
+
+rsloop_run.__sys.version_info = (3, 11)
+
+async def main():
+    try:
+        await asyncio.sleep(60)
+    finally:
+        print("main-cancelled", flush=True)
+
+threading.Timer(0.1, lambda: signal.raise_signal(signal.SIGINT)).start()
+try:
+    rsloop.run(main())
+except KeyboardInterrupt:
+    print("keyboard-interrupt", flush=True)
+    raise SystemExit(0)
+except BaseException as exc:
+    print(f"unexpected: {type(exc).__name__}: {exc}", flush=True)
+    raise SystemExit(2)
+else:
+    print("no-interrupt", flush=True)
+    raise SystemExit(3)
+"""
+        proc = subprocess.run(
+            [sys.executable, "-c", script],
+            text=True,
+            capture_output=True,
+            timeout=5,
+        )
+        self.assertEqual(proc.returncode, 0, proc.stdout + proc.stderr)
+        self.assertIn("main-cancelled", proc.stdout)
+        self.assertIn("keyboard-interrupt", proc.stdout)
+
     @unittest.skipUnless(
         os.name == "posix" and os.path.isdir("/proc/self/fd"),
         "requires /proc/self/fd",