Enable virtual terminal mode in Windows REPL

y5c4l3 · y5c4l3 · commit 718110f21181 · 2024-09-21T02:07:14.000+08:00
Windows REPL input was using virtual key mode, which does not support
terminal escape sequences. This patch calls `SetConsoleMode` properly
when initializing and send sequences to enable bracketed-paste modes
to support verbatim copy-and-paste.

Signed-off-by: y5c4l3 &lt;y5c4l3@proton.me&gt;
diff --git a/Lib/_pyrepl/windows_console.py b/Lib/_pyrepl/windows_console.py
@@ -42,6 +42,7 @@
 from .console import Event, Console
 from .trace import trace
 from .utils import wlen
+from .windows_eventqueue import EventQueue
 
 try:
     from ctypes import GetLastError, WinDLL, windll, WinError  # type: ignore[attr-defined]
@@ -94,7 +95,9 @@ def __init__(self, err: int | None, descr: str | None = None) -> None:
     0x83: "f20",  # VK_F20
 }
 
-# Console escape codes: https://learn.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences
+# Virtual terminal output sequences
+# Reference: https://learn.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences#output-sequences
+# Check `windows_eventqueue.py` for input sequences
 ERASE_IN_LINE = "\x1b[K"
 MOVE_LEFT = "\x1b[{}D"
 MOVE_RIGHT = "\x1b[{}C"
@@ -106,6 +109,12 @@ def __init__(self, err: int | None, descr: str | None = None) -> None:
 class _error(Exception):
     pass
 
+def _supports_vt():
+    try:
+        import nt
+        return nt._supports_virtual_terminal()
+    except (ImportError, AttributeError):
+        return False
 
 class WindowsConsole(Console):
     def __init__(
@@ -117,17 +126,36 @@ def __init__(
     ):
         super().__init__(f_in, f_out, term, encoding)
 
+        self.__vt_support = _supports_vt()
+        self.__vt_bracketed_paste = False
+
+        if self.__vt_support:
+            trace('console supports virtual terminal')
+
+            # Should make educated guess to determine the terminal type.
+            # Currently enable bracketed-paste only if it's Windows Terminal.
+            if 'WT_SESSION' in os.environ:
+                trace('console supports bracketed-paste sequence')
+                self.__vt_bracketed_paste = True
+
+        # Save original console modes so we can recover on cleanup.
+        original_input_mode = DWORD()
+        GetConsoleMode(InHandle, original_input_mode)
+        trace(f'saved original input mode 0x{original_input_mode.value:x}')
+        self.__original_input_mode = original_input_mode.value
+
         SetConsoleMode(
             OutHandle,
             ENABLE_WRAP_AT_EOL_OUTPUT
             | ENABLE_PROCESSED_OUTPUT
             | ENABLE_VIRTUAL_TERMINAL_PROCESSING,
         )
+
         self.screen: list[str] = []
         self.width = 80
         self.height = 25
         self.__offset = 0
-        self.event_queue: deque[Event] = deque()
+        self.event_queue = EventQueue(encoding)
         try:
             self.out = io._WindowsConsoleIO(self.output_fd, "w")  # type: ignore[attr-defined]
         except ValueError:
@@ -291,6 +319,12 @@ def _enable_blinking(self):
     def _disable_blinking(self):
         self.__write("\x1b[?12l")
 
+    def _enable_bracketed_paste(self) -> None:
+        self.__write("\x1b[?2004h")
+
+    def _disable_bracketed_paste(self) -> None:
+        self.__write("\x1b[?2004l")
+
     def __write(self, text: str) -> None:
         if "\x1a" in text:
             text = ''.join(["^Z" if x == '\x1a' else x for x in text])
@@ -320,8 +354,17 @@ def prepare(self) -> None:
         self.__gone_tall = 0
         self.__offset = 0
 
+        if self.__vt_support:
+            SetConsoleMode(InHandle, self.__original_input_mode | ENABLE_VIRTUAL_TERMINAL_INPUT)
+            if self.__vt_bracketed_paste:
+                self._enable_bracketed_paste()
+
     def restore(self) -> None:
-        pass
+        if self.__vt_support:
+            # Recover to original mode before running REPL
+            SetConsoleMode(InHandle, self.__original_input_mode)
+            if self.__vt_bracketed_paste:
+                self._disable_bracketed_paste()
 
     def _move_relative(self, x: int, y: int) -> None:
         """Moves relative to the current __posxy"""
@@ -342,7 +385,7 @@ def move_cursor(self, x: int, y: int) -> None:
             raise ValueError(f"Bad cursor position {x}, {y}")
 
         if y < self.__offset or y >= self.__offset + self.height:
-            self.event_queue.insert(0, Event("scroll", ""))
+            self.event_queue.insert(Event("scroll", ""))
         else:
             self._move_relative(x, y)
             self.__posxy = x, y
@@ -386,10 +429,8 @@ def get_event(self, block: bool = True) -> Event | None:
         """Return an Event instance.  Returns None if |block| is false
         and there is no event pending, otherwise waits for the
         completion of an event."""
-        if self.event_queue:
-            return self.event_queue.pop()
 
-        while True:
+        while self.event_queue.empty():
             rec = self._read_input()
             if rec is None:
                 if block:
@@ -428,8 +469,13 @@ def get_event(self, block: bool = True) -> Event | None:
                     continue
 
                 return None
+            elif self.__vt_support:
+                # If virtual terminal is enabled, scanning VT sequences
+                self.event_queue.push(rec.Event.KeyEvent.uChar.UnicodeChar)
+                continue
 
             return Event(evt="key", data=key, raw=rec.Event.KeyEvent.uChar.UnicodeChar)
+        return self.event_queue.get()
 
     def push_char(self, char: int | bytes) -> None:
         """
@@ -551,6 +597,13 @@ class INPUT_RECORD(Structure):
 MOUSE_EVENT = 0x02
 WINDOW_BUFFER_SIZE_EVENT = 0x04
 
+ENABLE_PROCESSED_INPUT = 0x0001
+ENABLE_LINE_INPUT = 0x0002
+ENABLE_ECHO_INPUT = 0x0004
+ENABLE_MOUSE_INPUT = 0x0010
+ENABLE_INSERT_MODE = 0x0020
+ENABLE_VIRTUAL_TERMINAL_INPUT = 0x0200
+
 ENABLE_PROCESSED_OUTPUT = 0x01
 ENABLE_WRAP_AT_EOL_OUTPUT = 0x02
 ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x04
@@ -582,6 +635,10 @@ class INPUT_RECORD(Structure):
     ]
     ScrollConsoleScreenBuffer.restype = BOOL
 
+    GetConsoleMode = _KERNEL32.GetConsoleMode
+    GetConsoleMode.argtypes = [HANDLE, POINTER(DWORD)]
+    GetConsoleMode.restype = BOOL
+
     SetConsoleMode = _KERNEL32.SetConsoleMode
     SetConsoleMode.argtypes = [HANDLE, DWORD]
     SetConsoleMode.restype = BOOL
@@ -600,6 +657,7 @@ def _win_only(*args, **kwargs):
     GetStdHandle = _win_only
     GetConsoleScreenBufferInfo = _win_only
     ScrollConsoleScreenBuffer = _win_only
+    GetConsoleMode = _win_only
     SetConsoleMode = _win_only
     ReadConsoleInput = _win_only
     OutHandle = 0
