fix(wizard): read keypresses from the raw fd so arrow keys work

sys.stdin.read(1) on the buffered TextIOWrapper pulls the whole escape
sequence into its Python-level buffer, so the follow-up select() sees an
empty fd and every arrow/Home/End press degrades to bare Esc; the unread
'[B' tail then leaks into the next keypress. Read raw bytes via os.read()
and select() on the file descriptor instead, keeping the same key tokens
and classifications. Add pty-driven regression tests over a real cbreak
terminal covering CSI sequences, single-byte keys, bare Esc, and tail-leak.
This commit is contained in:
2026-08-27 08:07:33 +00:00
parent 85ddfb4a71
commit 12ca06ba0d
2 changed files with 161 additions and 10 deletions
+13 -10
View File
@@ -343,23 +343,26 @@ def _read_keypress() -> str:
"""Read one keystroke from stdin (cbreak mode); classify it."""
import select
char = sys.stdin.read(1)
if char == "\x1b":
fd = sys.stdin.fileno()
# Read raw bytes from the fd: a buffered text read would pull the whole
# escape sequence into its own buffer, leaving nothing for select to see.
first = os.read(fd, 1)
if first == b"\x1b":
# A bare Esc is not followed by more bytes; CSI sequences are. Peek briefly.
readable, _, _ = select.select([sys.stdin], [], [], 0.05)
readable, _, _ = select.select([fd], [], [], 0.05)
if not readable:
return _KEY_ESCAPE
tail = sys.stdin.read(2)
return _CSI_KEYS.get(tail, _KEY_OTHER)
if char in ("\r", "\n"):
tail = os.read(fd, 2)
return _CSI_KEYS.get(tail.decode("ascii", "replace"), _KEY_OTHER)
if first in (b"\r", b"\n"):
return _KEY_ENTER
if char == "\x01": # Ctrl-A
if first == b"\x01": # Ctrl-A
return _KEY_HOME
if char == "\x05": # Ctrl-E
if first == b"\x05": # Ctrl-E
return _KEY_END
if char == "k":
if first == b"k":
return _KEY_UP
if char == "j":
if first == b"j":
return _KEY_DOWN
return _KEY_OTHER
+148
View File
@@ -2,8 +2,16 @@
from __future__ import annotations
from collections.abc import Callable
import io
import os
from pathlib import Path
import pty
import re
import select
import subprocess
import sys
import time
import unittest
from unittest import mock
@@ -54,5 +62,145 @@ class OptionLineRenderingTest(unittest.TestCase):
self.assertEqual(_option_line("Remove a source", False), " Remove a source")
_CHILD_SCRIPT = """
import os
import sys
import termios
import tty
sys.path.insert(0, os.environ["WIZARD_SRC"])
from weekly_activity.wizard import _read_keypress # noqa: E402
fd = sys.stdin.fileno()
saved = termios.tcgetattr(fd)
tty.setcbreak(fd)
try:
for _ in range(int(sys.argv[1])):
print("READY", flush=True)
print(f"KEY={_read_keypress()}", flush=True)
finally:
termios.tcsetattr(fd, termios.TCSADRAIN, saved)
"""
_KEY_LINE = re.compile(rb"KEY=([a-z]+)\r\n")
class ArrowKeyKeypressTest(unittest.TestCase):
"""CSI sequences must classify over a real pty instead of surfacing as bare Esc.
Drives the genuine ``_read_keypress()`` on a cbreak-mode terminal — the path the
arrow-key menu uses, unreachable through scripted ``io.StringIO`` stdin.
"""
TIMEOUT_S = 20.0
def _await(
self,
master: int,
child: subprocess.Popen[bytes],
box: dict[str, bytes],
condition: Callable[[], bool],
stage: str,
) -> None:
"""Block until ``condition`` holds on the pty stream, else fail loudly."""
deadline = time.monotonic() + self.TIMEOUT_S
while not condition():
remaining = deadline - time.monotonic()
if remaining <= 0:
break
readable, _, _ = select.select([master], [], [], min(remaining, 0.2))
if readable:
try:
chunk = os.read(master, 4096)
except OSError: # slave closed: the child exited early
break
if chunk:
box["stream"] += chunk
else:
return
if child.poll() is None:
child.kill()
child.wait(timeout=5)
err = child.stderr.read().decode("utf-8", "replace") if child.stderr else ""
self.fail(f"{stage}: stalled; tail={box['stream'][-300:]!r} stderr={err!r}")
def _child_tokens(self, sends: list[bytes]) -> list[str]:
"""Send each payload after its READY prompt; return printed key tokens."""
src_dir = str(Path(__file__).resolve().parents[1] / "src")
master, slave = pty.openpty()
child = subprocess.Popen(
[sys.executable, "-c", _CHILD_SCRIPT, str(len(sends))],
stdin=slave,
stdout=slave,
stderr=subprocess.PIPE,
env={**os.environ, "WIZARD_SRC": src_dir},
)
os.close(slave)
box: dict[str, bytes] = {"stream": b""}
try:
tokens: list[str] = []
for index, payload in enumerate(sends):
self._await(
master,
child,
box,
lambda i=index: box["stream"].count(b"READY") > i,
f"READY prompt {index} before {payload!r}",
)
seen = len(_KEY_LINE.findall(box["stream"]))
os.write(master, payload)
self._await(
master,
child,
box,
lambda n=seen: len(_KEY_LINE.findall(box["stream"])) > n,
f"token line after {payload!r}",
)
tokens.append(_KEY_LINE.findall(box["stream"])[-1].decode())
self.assertEqual(
box["stream"].count(b"READY"),
len(sends),
f"unexpected output tail {box['stream'][-300:]!r}",
)
self.assertEqual(len(_KEY_LINE.findall(box["stream"])), len(sends))
try:
child.wait(timeout=self.TIMEOUT_S)
except subprocess.TimeoutExpired:
child.kill()
child.wait()
self.fail("child ignored termination after completing every keystroke")
err = child.stderr.read().decode("utf-8", "replace") if child.stderr else ""
self.assertEqual(child.returncode, 0, err)
return tokens
finally:
if child.poll() is None:
child.kill()
child.wait()
if child.stderr is not None:
child.stderr.close()
os.close(master)
def test_arrow_home_end_sequences_map_to_selection_keys(self) -> None:
tokens = self._child_tokens([b"\x1b[A", b"\x1b[B", b"\x1b[H", b"\x1b[F"])
self.assertEqual(tokens, ["up", "down", "home", "end"])
def test_single_byte_keystrokes_still_classify(self) -> None:
# The reader switched to raw fd bytes; letters/Enter must classify as before.
tokens = self._child_tokens([b"\r", b"k", b"j", b"x"])
self.assertEqual(tokens, ["enter", "up", "down", "other"])
def test_bare_escape_is_classified_as_escape(self) -> None:
tokens = self._child_tokens([b"\x1b"])
self.assertEqual(tokens, ["escape"])
def test_escape_sequence_tail_does_not_leak_into_next_keypress(self) -> None:
# The original report: an arrow press poisoned the following reads.
tokens = self._child_tokens([b"\x1b[B", b"j", b"\r"])
self.assertEqual(tokens, ["down", "down", "enter"])
if __name__ == "__main__":
unittest.main()