diff --git a/src/invoke_toolkit/runners/rich.py b/src/invoke_toolkit/runners/rich.py index 05104c8..6c31dc9 100644 --- a/src/invoke_toolkit/runners/rich.py +++ b/src/invoke_toolkit/runners/rich.py @@ -3,15 +3,24 @@ in tasks. """ +import array +import errno import sys -from typing import TYPE_CHECKING +from typing import IO, TYPE_CHECKING -from invoke.runners import Local -from invoke.util import debug +from invoke.runners import Local, ready_for_reading +from invoke.util import debug, has_fileno, isatty from rich.syntax import Syntax from invoke_toolkit.output import get_console +if sys.platform != "win32": + import fcntl + import termios +else: + fcntl = None # pylint: disable=invalid-name + termios = None # pylint: disable=invalid-name + if TYPE_CHECKING: from invoke_toolkit.config.status_helper import StatusHelper @@ -96,9 +105,31 @@ def __getattr__(self, name): class NoStdoutRunner(Local): - """Invoke runner that prints to stderr when invoke is used with -e/--echo - and redacts secrets from subprocess output when redaction is enabled. - """ + """Invoke runner with output handling and safe POSIX stdin probing.""" + + def read_our_stdin(self, input_: IO) -> str | None: + """Read stdin without Invoke's undersized POSIX ioctl buffer.""" + if not ready_for_reading(input_): + return None + + bytes_to_read = 1 + if sys.platform != "win32" and isatty(input_) and has_fileno(input_): + buffer = array.array("i", [0]) + try: + fcntl.ioctl(input_, termios.FIONREAD, buffer, True) + bytes_to_read = max(buffer[0], 1) + except OSError: + pass + + try: + bytes_ = input_.read(bytes_to_read) + except OSError as e: + if e.errno != errno.EBADF: + raise + bytes_ = None + if bytes_ and isinstance(bytes_, bytes): + bytes_ = self.decode(bytes_) + return bytes_ def echo(self, command): if hasattr(self.context, "print"): diff --git a/tests/test_pty_stdin.py b/tests/test_pty_stdin.py new file mode 100644 index 0000000..7138664 --- /dev/null +++ b/tests/test_pty_stdin.py @@ -0,0 +1,24 @@ +import os +import tty +from contextlib import closing + +import pytest + +from invoke_toolkit import Context +from invoke_toolkit.runners.rich import NoStdoutRunner + + +@pytest.mark.skipif(os.name == "nt", reason="PTYs are POSIX-only") +def test_read_our_stdin_handles_escape_sequence_from_pty(): + master_fd, slave_fd = os.openpty() + try: + tty.setcbreak(slave_fd) + os.write(master_fd, b"\x1b[A") + stream = os.fdopen(slave_fd, "rb", closefd=False) + with closing(stream): + runner = NoStdoutRunner(Context()) + runner.encoding = "utf-8" + assert runner.read_our_stdin(stream) == "\x1b[A" + finally: + os.close(master_fd) + os.close(slave_fd)