From 033510e11dff742d9626b9fd895925ac77f566f1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Langa?= <lukasz@langa.pl>
Date: Fri, 6 Sep 2024 21:28:29 +0200
Subject: [PATCH] gh-120221: Support KeyboardInterrupt in asyncio REPL
 (#123795)

This switches the main pyrepl event loop to always be non-blocking so that it
can listen to incoming interruptions from other threads.

This also resolves invalid display of exceptions from other threads
(gh-123178).

This also fixes freezes with pasting and an active input hook.
---
 Lib/_pyrepl/_threading_handler.py             | 74 +++++++++++++++++++
 Lib/_pyrepl/reader.py                         | 42 +++++++----
 Lib/_pyrepl/unix_console.py                   | 15 +++-
 Lib/_pyrepl/windows_console.py                |  2 +-
 Lib/asyncio/__main__.py                       | 10 +++
 Lib/test/test_pyrepl/support.py               |  4 +-
 Lib/test/test_repl.py                         |  5 +-
 ...-09-06-19-23-44.gh-issue-120221.giJEDT.rst |  2 +
 8 files changed, 133 insertions(+), 21 deletions(-)
 create mode 100644 Lib/_pyrepl/_threading_handler.py
 create mode 100644 Misc/NEWS.d/next/Core_and_Builtins/2024-09-06-19-23-44.gh-issue-120221.giJEDT.rst

diff --git a/Lib/_pyrepl/_threading_handler.py b/Lib/_pyrepl/_threading_handler.py
new file mode 100644
index 00000000000..82f5e8650a2
--- /dev/null
+++ b/Lib/_pyrepl/_threading_handler.py
@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import traceback
+
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from threading import Thread
+    from types import TracebackType
+    from typing import Protocol
+
+    class ExceptHookArgs(Protocol):
+        @property
+        def exc_type(self) -> type[BaseException]: ...
+        @property
+        def exc_value(self) -> BaseException | None: ...
+        @property
+        def exc_traceback(self) -> TracebackType | None: ...
+        @property
+        def thread(self) -> Thread | None: ...
+
+    class ShowExceptions(Protocol):
+        def __call__(self) -> int: ...
+        def add(self, s: str) -> None: ...
+
+    from .reader import Reader
+
+
+def install_threading_hook(reader: Reader) -> None:
+    import threading
+
+    @dataclass
+    class ExceptHookHandler:
+        lock: threading.Lock = field(default_factory=threading.Lock)
+        messages: list[str] = field(default_factory=list)
+
+        def show(self) -> int:
+            count = 0
+            with self.lock:
+                if not self.messages:
+                    return 0
+                reader.restore()
+                for tb in self.messages:
+                    count += 1
+                    if tb:
+                        print(tb)
+                self.messages.clear()
+                reader.scheduled_commands.append("ctrl-c")
+                reader.prepare()
+            return count
+
+        def add(self, s: str) -> None:
+            with self.lock:
+                self.messages.append(s)
+
+        def exception(self, args: ExceptHookArgs) -> None:
+            lines = traceback.format_exception(
+                args.exc_type,
+                args.exc_value,
+                args.exc_traceback,
+                colorize=reader.can_colorize,
+            )  # type: ignore[call-overload]
+            pre = f"\nException in {args.thread.name}:\n" if args.thread else "\n"
+            tb = pre + "".join(lines)
+            self.add(tb)
+
+        def __call__(self) -> int:
+            return self.show()
+
+
+    handler = ExceptHookHandler()
+    reader.threading_hook = handler
+    threading.excepthook = handler.exception
diff --git a/Lib/_pyrepl/reader.py b/Lib/_pyrepl/reader.py
index aa3f5fd283e..54bd1ea0222 100644
--- a/Lib/_pyrepl/reader.py
+++ b/Lib/_pyrepl/reader.py
@@ -36,8 +36,7 @@ from .trace import trace
 
 # types
 Command = commands.Command
-if False:
-    from .types import Callback, SimpleContextManager, KeySpec, CommandName
+from .types import Callback, SimpleContextManager, KeySpec, CommandName
 
 
 def disp_str(buffer: str) -> tuple[str, list[int]]:
@@ -247,6 +246,7 @@ class Reader:
     lxy: tuple[int, int] = field(init=False)
     scheduled_commands: list[str] = field(default_factory=list)
     can_colorize: bool = False
+    threading_hook: Callback | None = None
 
     ## cached metadata to speed up screen refreshes
     @dataclass
@@ -722,6 +722,24 @@ class Reader:
             self.console.finish()
             self.finish()
 
+    def run_hooks(self) -> None:
+        threading_hook = self.threading_hook
+        if threading_hook is None and 'threading' in sys.modules:
+            from ._threading_handler import install_threading_hook
+            install_threading_hook(self)
+        if threading_hook is not None:
+            try:
+                threading_hook()
+            except Exception:
+                pass
+
+        input_hook = self.console.input_hook
+        if input_hook:
+            try:
+                input_hook()
+            except Exception:
+                pass
+
     def handle1(self, block: bool = True) -> bool:
         """Handle a single event.  Wait as long as it takes if block
         is true (the default), otherwise return False if no event is
@@ -732,16 +750,13 @@ class Reader:
             self.dirty = True
 
         while True:
-            input_hook = self.console.input_hook
-            if input_hook:
-                input_hook()
-                # We use the same timeout as in readline.c: 100ms
-                while not self.console.wait(100):
-                    input_hook()
-                event = self.console.get_event(block=False)
-            else:
-                event = self.console.get_event(block)
-            if not event:  # can only happen if we're not blocking
+            # We use the same timeout as in readline.c: 100ms
+            self.run_hooks()
+            self.console.wait(100)
+            event = self.console.get_event(block=False)
+            if not event:
+                if block:
+                    continue
                 return False
 
             translate = True
@@ -763,8 +778,7 @@ class Reader:
             if cmd is None:
                 if block:
                     continue
-                else:
-                    return False
+                return False
 
             self.do_cmd(cmd)
             return True
diff --git a/Lib/_pyrepl/unix_console.py b/Lib/_pyrepl/unix_console.py
index 2f150371297..2576b938a34 100644
--- a/Lib/_pyrepl/unix_console.py
+++ b/Lib/_pyrepl/unix_console.py
@@ -199,8 +199,14 @@ class UnixConsole(Console):
         self.event_queue = EventQueue(self.input_fd, self.encoding)
         self.cursor_visible = 1
 
+    def more_in_buffer(self) -> bool:
+        return bool(
+            self.input_buffer
+            and self.input_buffer_pos < len(self.input_buffer)
+        )
+
     def __read(self, n: int) -> bytes:
-        if not self.input_buffer or self.input_buffer_pos >= len(self.input_buffer):
+        if not self.more_in_buffer():
             self.input_buffer = os.read(self.input_fd, 10000)
 
         ret = self.input_buffer[self.input_buffer_pos : self.input_buffer_pos + n]
@@ -393,6 +399,7 @@ class UnixConsole(Console):
         """
         if not block and not self.wait(timeout=0):
             return None
+
         while self.event_queue.empty():
             while True:
                 try:
@@ -413,7 +420,11 @@ class UnixConsole(Console):
         """
         Wait for events on the console.
         """
-        return bool(self.pollob.poll(timeout))
+        return (
+            not self.event_queue.empty()
+            or self.more_in_buffer()
+            or bool(self.pollob.poll(timeout))
+        )
 
     def set_cursor_vis(self, visible):
         """
diff --git a/Lib/_pyrepl/windows_console.py b/Lib/_pyrepl/windows_console.py
index 08337af8e7b..f7a0095d795 100644
--- a/Lib/_pyrepl/windows_console.py
+++ b/Lib/_pyrepl/windows_console.py
@@ -479,7 +479,7 @@ class WindowsConsole(Console):
         while True:
             if msvcrt.kbhit(): # type: ignore[attr-defined]
                 return True
-            if timeout and time.time() - start_time > timeout:
+            if timeout and time.time() - start_time > timeout / 1000:
                 return False
             time.sleep(0.01)
 
diff --git a/Lib/asyncio/__main__.py b/Lib/asyncio/__main__.py
index 111b7d92367..5120140e061 100644
--- a/Lib/asyncio/__main__.py
+++ b/Lib/asyncio/__main__.py
@@ -127,6 +127,15 @@ class REPLThread(threading.Thread):
 
             loop.call_soon_threadsafe(loop.stop)
 
+    def interrupt(self) -> None:
+        if not CAN_USE_PYREPL:
+            return
+
+        from _pyrepl.simple_interact import _get_reader
+        r = _get_reader()
+        if r.threading_hook is not None:
+            r.threading_hook.add("")  # type: ignore
+
 
 if __name__ == '__main__':
     sys.audit("cpython.run_stdin")
@@ -184,6 +193,7 @@ if __name__ == '__main__':
             keyboard_interrupted = True
             if repl_future and not repl_future.done():
                 repl_future.cancel()
+            repl_thread.interrupt()
             continue
         else:
             break
diff --git a/Lib/test/test_pyrepl/support.py b/Lib/test/test_pyrepl/support.py
index cb5cb4ab20a..672d4896c92 100644
--- a/Lib/test/test_pyrepl/support.py
+++ b/Lib/test/test_pyrepl/support.py
@@ -161,8 +161,8 @@ class FakeConsole(Console):
     def forgetinput(self) -> None:
         pass
 
-    def wait(self) -> None:
-        pass
+    def wait(self, timeout: float | None = None) -> bool:
+        return True
 
     def repaint(self) -> None:
         pass
diff --git a/Lib/test/test_repl.py b/Lib/test/test_repl.py
index cd8ef0f1057..7a7285a1a2f 100644
--- a/Lib/test/test_repl.py
+++ b/Lib/test/test_repl.py
@@ -242,6 +242,7 @@ class TestInteractiveInterpreter(unittest.TestCase):
     def test_asyncio_repl_is_ok(self):
         m, s = pty.openpty()
         cmd = [sys.executable, "-I", "-m", "asyncio"]
+        env = os.environ.copy()
         proc = subprocess.Popen(
             cmd,
             stdin=s,
@@ -249,7 +250,7 @@ class TestInteractiveInterpreter(unittest.TestCase):
             stderr=s,
             text=True,
             close_fds=True,
-            env=os.environ,
+            env=env,
         )
         os.close(s)
         os.write(m, b"await asyncio.sleep(0)\n")
@@ -270,7 +271,7 @@ class TestInteractiveInterpreter(unittest.TestCase):
             proc.kill()
             exit_code = proc.wait()
 
-        self.assertEqual(exit_code, 0)
+        self.assertEqual(exit_code, 0, "".join(output))
 
 class TestInteractiveModeSyntaxErrors(unittest.TestCase):
 
diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2024-09-06-19-23-44.gh-issue-120221.giJEDT.rst b/Misc/NEWS.d/next/Core_and_Builtins/2024-09-06-19-23-44.gh-issue-120221.giJEDT.rst
new file mode 100644
index 00000000000..c562b87b02a
--- /dev/null
+++ b/Misc/NEWS.d/next/Core_and_Builtins/2024-09-06-19-23-44.gh-issue-120221.giJEDT.rst
@@ -0,0 +1,2 @@
+asyncio REPL is now again properly recognizing KeyboardInterrupts. Display
+of exceptions raised in secondary threads is fixed.