grab token and install handler in one go

Author: richardsheridan

src/trio/_repl.py

@@ -4,9 +4,10 @@
 import contextlib
 import inspect
 import sys
-import types
 import warnings
 from code import InteractiveConsole
+from types import CodeType, FunctionType, FrameType
+from typing import Callable
 
 import outcome
 
@@ -19,6 +20,16 @@ class SuppressDecorator(contextlib.ContextDecorator, contextlib.suppress):
     pass
 
 
+@SuppressDecorator(KeyboardInterrupt)
+@trio.lowlevel.disable_ki_protection
+def terminal_newline() -> None:
+    import fcntl
+    import termios
+
+    # Fake up a newline char as if user had typed it at the terminal
+    fcntl.ioctl(sys.stdin, termios.TIOCSTI, b"\n")
+
+
 @final
 class TrioInteractiveConsole(InteractiveConsole):
     # code.InteractiveInterpreter defines locals as Mapping[str, Any]
@@ -31,8 +42,8 @@ def __init__(self, repl_locals: dict[str, object] | None = None) -> None:
         self.token: trio.lowlevel.TrioToken | None = None
         self.compile.compiler.flags |= ast.PyCF_ALLOW_TOP_LEVEL_AWAIT
 
-    def runcode(self, code: types.CodeType) -> None:
-        func = types.FunctionType(code, self.locals)
+    def runcode(self, code: CodeType) -> None:
+        func = FunctionType(code, self.locals)
         if inspect.iscoroutinefunction(func):
             result = trio.from_thread.run(outcome.acapture, func)
         else:
@@ -79,24 +90,19 @@ def raw_input(self, prompt: str = "") -> str:
 
             interrupted = False
 
-            if self.token is None:
-                self.token = trio.from_thread.run_sync(trio.lowlevel.current_trio_token)
-
-            @SuppressDecorator(KeyboardInterrupt)
-            @trio.lowlevel.disable_ki_protection
-            def newline():
-                import fcntl
-                import termios
+            def install_handler() -> (
+                Callable[[int, FrameType | None], None] | int | None
+            ):
+                def handler(sig: int, frame: FrameType | None) -> None:
+                    nonlocal interrupted
+                    interrupted = True
+                    token.run_sync_soon(terminal_newline, idempotent=True)
 
-                # Fake up a newline char as if user had typed it at
-                fcntl.ioctl(sys.stdin, termios.TIOCSTI, b"\n")
+                token = trio.lowlevel.current_trio_token()
 
-            def handler(sig: int, frame: types.FrameType | None) -> None:
-                nonlocal interrupted
-                interrupted = True
-                self.token.run_sync_soon(newline, idempotent=True)
+                return signal(SIGINT, handler)
 
-            prev_handler = trio.from_thread.run_sync(signal, SIGINT, handler)
+            prev_handler = trio.from_thread.run_sync(install_handler)
             try:
                 return input(prompt)
             finally: