Properly detect single ESC key

Author: lkollar

Lib/_pyrepl/base_eventqueue.py

@@ -31,13 +31,16 @@
 from .trace import trace
 
 class BaseEventQueue:
+    _ESCAPE_TIMEOUT_MS = 50
+
     def __init__(self, encoding: str, keymap_dict: dict[bytes, str]) -> None:
         self.compiled_keymap = keymap.compile_keymap(keymap_dict)
         self.keymap = self.compiled_keymap
         trace("keymap {k!r}", k=self.keymap)
         self.encoding = encoding
         self.events: deque[Event] = deque()
         self.buf = bytearray()
+        self._pending_escape_deadline: float | None = None
 
     def get(self) -> Event | None:
         """
@@ -69,6 +72,50 @@ def insert(self, event: Event) -> None:
         trace('added event {event}', event=event)
         self.events.append(event)
 
+    def has_pending_escape_sequence(self) -> bool:
+        """
+        Check if there's a potential escape sequence waiting for more input.
+
+        Returns True if we have exactly one byte (ESC) in the buffer and
+        we're in the middle of keymap navigation, indicating we're waiting
+        to see if more bytes will arrive to complete an escape sequence.
+        """
+        return (
+            len(self.buf) == 1
+            and self.buf[0] == 27  # ESC byte
+            and self.keymap is not self.compiled_keymap
+        )
+
+    def should_emit_standalone_escape(self, current_time_ms: float) -> bool:
+        """
+        Check if a pending ESC should be emitted as a standalone escape key.
+        """
+        if not self.has_pending_escape_sequence():
+            return False
+
+        if self._pending_escape_deadline is None:
+            # First time checking - set the deadline
+            self._pending_escape_deadline = current_time_ms + self._ESCAPE_TIMEOUT_MS
+            return False
+
+        # Check if the deadline has passed
+        return current_time_ms >= self._pending_escape_deadline
+
+    def emit_standalone_escape(self) -> None:
+        """
+        Emit the buffered ESC byte as a standalone escape key event.
+        """
+        self.keymap = self.compiled_keymap
+        # Standalone ESC event
+        self.insert(Event('key', '\033', b'\033'))
+
+        # Just in case there are remaining bytes in the buffer
+        remaining = self.flush_buf()[1:]
+        for byte in remaining:
+            self.push(byte)
+
+        self._pending_escape_deadline = None
+
     def push(self, char: int | bytes) -> None:
         """
         Processes a character by updating the buffer and handling special key mappings.
@@ -78,6 +125,9 @@ def push(self, char: int | bytes) -> None:
         char = ord_char.to_bytes()
         self.buf.append(ord_char)
 
+        if self._pending_escape_deadline is not None:
+            self._pending_escape_deadline = None
+
         if char in self.keymap:
             if self.keymap is self.compiled_keymap:
                 # sanity check, buffer is empty when a special key comes

Lib/_pyrepl/unix_console.py

@@ -419,6 +419,21 @@ def get_event(self, block: bool = True) -> Event | None:
             return None
 
         while self.event_queue.empty():
+            # Check if we have a pending escape sequence that needs timeout handling
+            if self.event_queue.has_pending_escape_sequence():
+                current_time_ms = time.monotonic() * 1000
+
+                if self.event_queue.should_emit_standalone_escape(current_time_ms):
+                    # Timeout expired - emit the ESC as a standalone key
+                    self.event_queue.emit_standalone_escape()
+                    break
+
+                if not self.wait(timeout=10):
+                    current_time_ms = time.monotonic() * 1000
+                    if self.event_queue.should_emit_standalone_escape(current_time_ms):
+                        self.event_queue.emit_standalone_escape()
+                    continue
+
             while True:
                 try:
                     self.push_char(self.__read(1))
@@ -445,6 +460,7 @@ def wait(self, timeout: float | None = None) -> bool:
             or bool(self.pollob.poll(timeout))
         )
 
+
     def set_cursor_vis(self, visible):
         """
         Set the visibility of the cursor.

Lib/_pyrepl/windows_console.py

@@ -446,6 +446,24 @@ def get_event(self, block: bool = True) -> Event | None:
             return None
 
         while self.event_queue.empty():
+            # Check if we have a pending escape sequence that needs timeout handling
+            if self.event_queue.has_pending_escape_sequence():
+                import time
+                current_time_ms = time.monotonic() * 1000
+
+                if self.event_queue.should_emit_standalone_escape(current_time_ms):
+                    # Timeout expired - emit the ESC as a standalone key
+                    self.event_queue.emit_standalone_escape()
+                    break
+
+                # Wait for a short time to check for more input
+                if not self.wait(timeout=10):
+                    # Check again after timeout
+                    current_time_ms = time.monotonic() * 1000
+                    if self.event_queue.should_emit_standalone_escape(current_time_ms):
+                        self.event_queue.emit_standalone_escape()
+                    continue
+
             rec = self._read_input()
             if rec is None:
                 return None
@@ -583,6 +601,7 @@ def repaint(self) -> None:
         raise NotImplementedError("No repaint support")
 
 
+
 # Windows interop
 class CONSOLE_SCREEN_BUFFER_INFO(Structure):
     _fields_ = [

Lib/test/test_pyrepl/test_reader.py

@@ -551,6 +551,22 @@ def test_syntax_highlighting_literal_brace_in_fstring_or_tstring(self):
         self.maxDiff=None
         self.assert_screen_equal(reader, expected)
 
+    def test_vi_escape_switches_to_normal_mode(self):
+        events = itertools.chain(
+            code_to_events("hello"),
+            [
+                Event(evt="key", data="\x1b", raw=bytearray(b"\x1b")),
+                Event(evt="key", data="h", raw=bytearray(b"h")),
+            ],
+        )
+        reader, _ = handle_all_events(
+            events,
+            prepare_reader=prepare_vi_reader,
+        )
+        self.assertEqual(reader.get_unicode(), "hello")
+        self.assertTrue(reader.editor_mode.is_normal())
+        self.assertEqual(reader.pos, len("hello") - 2)  # After 'h' left movement
+
     def test_control_characters(self):
         code = 'flag = "🏳️‍🌈"'
         events = code_to_events(code)