Skip to content

Commit 0fed3f7

Browse files
committed
Enable virtual terminal mode in Windows REPL
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 <[email protected]>
1 parent 6d96558 commit 0fed3f7

File tree

1 file changed

+65
-7
lines changed

1 file changed

+65
-7
lines changed

Lib/_pyrepl/windows_console.py

Lines changed: 65 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
from .console import Event, Console
4343
from .trace import trace
4444
from .utils import wlen
45+
from .windows_eventqueue import EventQueue
4546

4647
try:
4748
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:
9495
0x83: "f20", # VK_F20
9596
}
9697

97-
# Console escape codes: https://learn.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences
98+
# Virtual terminal output sequences
99+
# Reference: https://learn.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences#output-sequences
100+
# Check `windows_eventqueue.py` for input sequences
98101
ERASE_IN_LINE = "\x1b[K"
99102
MOVE_LEFT = "\x1b[{}D"
100103
MOVE_RIGHT = "\x1b[{}C"
@@ -106,6 +109,12 @@ def __init__(self, err: int | None, descr: str | None = None) -> None:
106109
class _error(Exception):
107110
pass
108111

112+
def _supports_vt():
113+
try:
114+
import nt
115+
return nt._supports_virtual_terminal()
116+
except (ImportError, AttributeError):
117+
return False
109118

110119
class WindowsConsole(Console):
111120
def __init__(
@@ -117,17 +126,36 @@ def __init__(
117126
):
118127
super().__init__(f_in, f_out, term, encoding)
119128

129+
self.__vt_support = _supports_vt()
130+
self.__vt_bracketed_paste = False
131+
132+
if self.__vt_support:
133+
trace('console supports virtual terminal')
134+
135+
# Should make educated guess to determine the terminal type.
136+
# Currently enable bracketed-paste only if it's Windows Terminal.
137+
if 'WT_SESSION' in os.environ:
138+
trace('console supports bracketed-paste sequence')
139+
self.__vt_bracketed_paste = True
140+
141+
# Save original console modes so we can recover on cleanup.
142+
original_input_mode = DWORD()
143+
GetConsoleMode(InHandle, original_input_mode)
144+
trace(f'saved original input mode 0x{original_input_mode.value:x}')
145+
self.__original_input_mode = original_input_mode.value
146+
120147
SetConsoleMode(
121148
OutHandle,
122149
ENABLE_WRAP_AT_EOL_OUTPUT
123150
| ENABLE_PROCESSED_OUTPUT
124151
| ENABLE_VIRTUAL_TERMINAL_PROCESSING,
125152
)
153+
126154
self.screen: list[str] = []
127155
self.width = 80
128156
self.height = 25
129157
self.__offset = 0
130-
self.event_queue: deque[Event] = deque()
158+
self.event_queue = EventQueue(encoding)
131159
try:
132160
self.out = io._WindowsConsoleIO(self.output_fd, "w") # type: ignore[attr-defined]
133161
except ValueError:
@@ -291,6 +319,12 @@ def _enable_blinking(self):
291319
def _disable_blinking(self):
292320
self.__write("\x1b[?12l")
293321

322+
def _enable_bracketed_paste(self) -> None:
323+
self.__write("\x1b[?2004h")
324+
325+
def _disable_bracketed_paste(self) -> None:
326+
self.__write("\x1b[?2004l")
327+
294328
def __write(self, text: str) -> None:
295329
if "\x1a" in text:
296330
text = ''.join(["^Z" if x == '\x1a' else x for x in text])
@@ -320,8 +354,17 @@ def prepare(self) -> None:
320354
self.__gone_tall = 0
321355
self.__offset = 0
322356

357+
if self.__vt_support:
358+
SetConsoleMode(InHandle, self.__original_input_mode | ENABLE_VIRTUAL_TERMINAL_INPUT)
359+
if self.__vt_bracketed_paste:
360+
self._enable_bracketed_paste()
361+
323362
def restore(self) -> None:
324-
pass
363+
if self.__vt_support:
364+
# Recover to original mode before running REPL
365+
SetConsoleMode(InHandle, self.__original_input_mode)
366+
if self.__vt_bracketed_paste:
367+
self._disable_bracketed_paste()
325368

326369
def _move_relative(self, x: int, y: int) -> None:
327370
"""Moves relative to the current __posxy"""
@@ -342,7 +385,7 @@ def move_cursor(self, x: int, y: int) -> None:
342385
raise ValueError(f"Bad cursor position {x}, {y}")
343386

344387
if y < self.__offset or y >= self.__offset + self.height:
345-
self.event_queue.insert(0, Event("scroll", ""))
388+
self.event_queue.insert(Event("scroll", ""))
346389
else:
347390
self._move_relative(x, y)
348391
self.__posxy = x, y
@@ -390,10 +433,8 @@ def get_event(self, block: bool = True) -> Event | None:
390433
"""Return an Event instance. Returns None if |block| is false
391434
and there is no event pending, otherwise waits for the
392435
completion of an event."""
393-
if self.event_queue:
394-
return self.event_queue.pop()
395436

396-
while True:
437+
while self.event_queue.empty():
397438
rec = self._read_input(block)
398439
if rec is None:
399440
return None
@@ -430,8 +471,13 @@ def get_event(self, block: bool = True) -> Event | None:
430471
continue
431472

432473
return None
474+
elif self.__vt_support:
475+
# If virtual terminal is enabled, scanning VT sequences
476+
self.event_queue.push(rec.Event.KeyEvent.uChar.UnicodeChar)
477+
continue
433478

434479
return Event(evt="key", data=key, raw=rec.Event.KeyEvent.uChar.UnicodeChar)
480+
return self.event_queue.get()
435481

436482
def push_char(self, char: int | bytes) -> None:
437483
"""
@@ -553,6 +599,13 @@ class INPUT_RECORD(Structure):
553599
MOUSE_EVENT = 0x02
554600
WINDOW_BUFFER_SIZE_EVENT = 0x04
555601

602+
ENABLE_PROCESSED_INPUT = 0x0001
603+
ENABLE_LINE_INPUT = 0x0002
604+
ENABLE_ECHO_INPUT = 0x0004
605+
ENABLE_MOUSE_INPUT = 0x0010
606+
ENABLE_INSERT_MODE = 0x0020
607+
ENABLE_VIRTUAL_TERMINAL_INPUT = 0x0200
608+
556609
ENABLE_PROCESSED_OUTPUT = 0x01
557610
ENABLE_WRAP_AT_EOL_OUTPUT = 0x02
558611
ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x04
@@ -584,6 +637,10 @@ class INPUT_RECORD(Structure):
584637
]
585638
ScrollConsoleScreenBuffer.restype = BOOL
586639

640+
GetConsoleMode = _KERNEL32.GetConsoleMode
641+
GetConsoleMode.argtypes = [HANDLE, POINTER(DWORD)]
642+
GetConsoleMode.restype = BOOL
643+
587644
SetConsoleMode = _KERNEL32.SetConsoleMode
588645
SetConsoleMode.argtypes = [HANDLE, DWORD]
589646
SetConsoleMode.restype = BOOL
@@ -610,6 +667,7 @@ def _win_only(*args, **kwargs):
610667
GetStdHandle = _win_only
611668
GetConsoleScreenBufferInfo = _win_only
612669
ScrollConsoleScreenBuffer = _win_only
670+
GetConsoleMode = _win_only
613671
SetConsoleMode = _win_only
614672
ReadConsoleInput = _win_only
615673
GetNumberOfConsoleInputEvents = _win_only

0 commit comments

Comments
 (0)