"""ShellNode: Node wrapping subprocess.Popen with non-blocking I/O.
Supports both pipe-based (default) and PTY-based subprocess management.
PTY mode is required for interactive terminal applications like bash, nano, vim.
Pipe mode works on every platform; PTY mode is POSIX only.
"""
import codecs
import logging
import os
import queue
import selectors
import shlex
import struct
import subprocess
import sys
import threading
import time
from typing import IO
from .descriptors import Property
from .node import Node
from .signals import Signal
log = logging.getLogger(__name__)
_READ_SIZE = 65536
# Incremental, so a UTF-8 character split across two reads is decoded whole.
_utf8_decoder = codecs.getincrementaldecoder("utf-8")
if sys.platform == "win32":
import _winapi
import ctypes
import msvcrt
from ctypes import wintypes
_kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
_OpenThread = _kernel32.OpenThread
_OpenThread.argtypes = (wintypes.DWORD, wintypes.BOOL, wintypes.DWORD)
_OpenThread.restype = wintypes.HANDLE
_CancelSynchronousIo = _kernel32.CancelSynchronousIo
_CancelSynchronousIo.argtypes = (wintypes.HANDLE,)
_CancelSynchronousIo.restype = wintypes.BOOL
_CloseHandle = _kernel32.CloseHandle
_CloseHandle.argtypes = (wintypes.HANDLE,)
_CloseHandle.restype = wintypes.BOOL
_THREAD_TERMINATE = 0x0001 # the access right CancelSynchronousIo requires
_ERROR_OPERATION_ABORTED = 995 # Windows: a read cancelled by CancelSynchronousIo
class _SelectorReader:
"""Reads the output streams on the frame thread through a zero-timeout selector: POSIX.
No threads and nothing running while the child is quiet. A stream at end of file
stays open until ``close()``, so a PTY master stays valid for writes.
"""
# Reads per stream once the child has exited. A pipe holds at most 1 MiB on Linux
# (``/proc/sys/fs/pipe-max-size``), so this empties it; the cap only stops a
# grandchild that keeps writing from holding the frame.
_EXIT_READS = 16
def __init__(self, streams: list[tuple[IO[bytes], str]]):
self._streams = [stream for stream, _ in streams]
self._sel = selectors.DefaultSelector()
for stream, name in streams:
os.set_blocking(stream.fileno(), False)
self._sel.register(stream, selectors.EVENT_READ, (name, _utf8_decoder(errors="replace")))
self.done = False
def read(self, exited: bool) -> list[tuple[str, str]]:
"""Return ``(stream name, text)`` for what arrived since the last call, without blocking.
While the child runs each ready stream is read once. Once it has exited, every stream is
read until it is empty or at end of file and the reader is ``done``: what a grandchild
holding a pipe writes later is not the child's output.
"""
out = []
for key, _ in self._sel.select(timeout=0):
name, decoder = key.data
for _ in range(self._EXIT_READS if exited else 1):
try:
chunk = os.read(key.fd, _READ_SIZE)
except BlockingIOError:
break
except OSError: # a PTY master reports EIO once the terminal closes
chunk = b""
if text := decoder.decode(chunk, final=not chunk):
out.append((name, text))
if not chunk:
self._sel.unregister(key.fileobj)
break
self.done = exited
return out
def close(self):
self._sel.close()
for stream in self._streams:
stream.close()
class _ThreadReader:
"""Reads each stream on its own thread, blocked in ``read`` until data or EOF: Windows.
Windows cannot select on a pipe, and a blocked read costs nothing while the child is
quiet. Each thread owns its stream and ends at end of file; ``close()`` cancels a read
still blocked because a grandchild holds the pipe open, so no thread outlives the node.
"""
# How long after the child exits a thread may still be delivering its output. The
# frame cannot see what a pipe holds, so a grandchild holding one open is waited out.
_EXIT_GRACE_SECONDS = 0.25
def __init__(self, streams: list[tuple[IO[bytes], str]]):
self._queue: queue.SimpleQueue[tuple[str, str]] = queue.SimpleQueue()
self._stopping = False
self._exited_at: float | None = None
self._threads = [
threading.Thread(target=self._pump, args=(stream, name), name=f"ShellNode-{name}", daemon=True)
for stream, name in streams
]
for thread in self._threads:
thread.start()
@staticmethod
def _read_blocking(fd: int) -> bytes:
"""Read up to ``_READ_SIZE`` bytes from the pipe ``fd``, waiting for data; ``b""`` at EOF.
On Windows through ReadFile, so a cancelled read raises with ``winerror`` set.
"""
if sys.platform == "win32":
try:
return _winapi.ReadFile(msvcrt.get_osfhandle(fd), _READ_SIZE)[0]
except BrokenPipeError: # every write end is closed
return b""
return os.read(fd, _READ_SIZE)
def _pump(self, stream: IO[bytes], name: str):
decoder = _utf8_decoder(errors="replace")
with stream:
while True:
try:
chunk = self._read_blocking(stream.fileno())
except OSError as e:
if not (self._stopping and getattr(e, "winerror", None) == _ERROR_OPERATION_ABORTED):
log.error("ShellNode %s read failed: %s", name, e)
chunk = b""
if text := decoder.decode(chunk, final=not chunk):
self._queue.put((name, text))
if not chunk:
return
def read(self, exited: bool) -> list[tuple[str, str]]:
if exited and self._exited_at is None:
self._exited_at = time.monotonic()
return [self._queue.get_nowait() for _ in range(self._queue.qsize())]
@property
def done(self) -> bool:
"""The child has exited and its output is all delivered, or the grace has run out."""
if self._exited_at is None:
return False
# Alive first: once every thread has ended, nothing more can reach the queue.
drained = not any(thread.is_alive() for thread in self._threads) and self._queue.empty()
return drained or time.monotonic() - self._exited_at >= self._EXIT_GRACE_SECONDS
def close(self):
"""Cancel the reads still blocked, so every thread ends now.
CancelSynchronousIo fails a pending synchronous ReadFile with ERROR_OPERATION_ABORTED
and does nothing while the thread is between reads, so it repeats until the thread ends:
https://learn.microsoft.com/en-us/windows/win32/fileio/cancelsynchronousio-func
POSIX, where only tests run this reader, has no such call, and a thread ends at EOF.
"""
self._stopping = True
if sys.platform == "win32":
for thread in self._threads:
assert thread.native_id is not None # set once the thread has started
if not (handle := _OpenThread(_THREAD_TERMINATE, False, thread.native_id)):
if thread.is_alive():
error = ctypes.get_last_error()
log.error("ShellNode cannot cancel %s: OpenThread failed (%d)", thread.name, error)
continue
while thread.is_alive():
_CancelSynchronousIo(handle)
thread.join(0.01)
_CloseHandle(handle)
# The pipe-mode reader for this platform; a test may swap it to run the other backend.
_PipeReader: type[_SelectorReader | _ThreadReader] = _ThreadReader if sys.platform == "win32" else _SelectorReader
[docs]
class ShellNode(Node):
"""Node that manages a subprocess with non-blocking I/O.
Emits stdout_data(str) and stderr_data(str) as output arrives, and
process_exited(int) when the child terminates. Output is delivered in
on_update and never blocks the game loop: on POSIX the frame reads the pipes
through a zero-timeout selector; on Windows a thread per pipe waits in a read.
``process_exited`` comes after all the child's output. A grandchild that
inherited the pipes (a background job, a worker an LSP server spawned) can
hold them open indefinitely, so the node does not wait for them to close: on
POSIX ``process_exited`` fires in the frame the child's exit is seen, and on
Windows after at most a short grace. The node then stops reading and leaves
the grandchild running. Stopping the node, or its leaving the tree, also
stops reading.
For interactive terminal applications, set use_pty=True to allocate
a pseudo-terminal. In PTY mode, stdout and stderr are merged. PTY mode is
POSIX only: on Windows ``start()`` raises ``NotImplementedError``.
"""
auto_start = Property(False, hint="Start process on tree entry")
command = Property("", hint="Command to execute")
def __init__(self, command: str = "", *, use_pty: bool = False, env: dict[str, str] | None = None, **kwargs):
super().__init__(**kwargs)
self.command = command
self.use_pty = use_pty
self.env = env
self.stdout_data = Signal()
self.stderr_data = Signal()
self.process_exited = Signal()
self._proc: subprocess.Popen | None = None
self._exit_code: int | None = None
self._reader: _SelectorReader | _ThreadReader | None = None
# PTY mode: the master side, which stdin writes and resize go to.
self._master_fd: int | None = None
[docs]
def on_enter_tree(self):
# Entry, not ready: ``on_exit_tree`` stops the process, so the restart
# owed to a node that leaves the tree and comes back belongs to the hook
# that runs on every entry. ``on_ready`` runs once per node, ever, and
# would leave a re-added shell dead.
if self.auto_start and self.command and self._proc is None:
self.start()
[docs]
def start(self, command: str | None = None):
"""Launch the subprocess. Replaces any already-running process."""
if self.use_pty and sys.platform == "win32":
raise NotImplementedError(
"ShellNode(use_pty=True) needs a POSIX pseudo-terminal, which Windows does not provide. "
"Use pipe mode (use_pty=False) on Windows; interactive terminal programs that need a TTY "
"are not supported there."
)
if command is not None:
self.command = command
if not self.command:
log.error("ShellNode.start() called with empty command")
return
if self._proc is not None:
self.stop()
self._exit_code = None
if self.use_pty:
self._start_pty(shlex.split(self.command))
else:
# Windows takes a command line string as it is; splitting it with POSIX
# rules would eat the backslashes in its paths.
self._start_pipe(self.command if sys.platform == "win32" else shlex.split(self.command))
def _start_pipe(self, args: list[str] | str):
"""Start subprocess with stdin/stdout/stderr pipes."""
try:
self._proc = subprocess.Popen(
args,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
env=self.env or os.environ.copy(),
)
except OSError as e:
log.error("Failed to start process %r: %s", self.command, e)
return
assert self._proc.stdout is not None and self._proc.stderr is not None
self._reader = _PipeReader([(self._proc.stdout, "stdout"), (self._proc.stderr, "stderr")])
def _start_pty(self, args: list[str]):
"""Start subprocess with a pseudo-terminal."""
import pty
master, slave = pty.openpty()
env = (self.env or os.environ).copy()
env["TERM"] = "xterm-256color"
try:
self._proc = subprocess.Popen(
args,
stdin=slave,
stdout=slave,
stderr=slave,
close_fds=True,
start_new_session=True,
env=env,
)
except OSError as e:
os.close(master)
os.close(slave)
log.error("Failed to start PTY process %r: %s", self.command, e)
return
os.close(slave)
self._master_fd = master
# The reader owns the master and closes it; writes and resizes use the same descriptor.
self._reader = _SelectorReader([(open(master, "rb", buffering=0), "stdout")])
[docs]
def resize(self, cols: int, rows: int):
"""Set the PTY terminal size. Only works in PTY mode."""
if self._master_fd is None:
return
import fcntl
import signal
import termios
fcntl.ioctl(self._master_fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
if self._proc and self._proc.pid:
try:
os.killpg(os.getpgid(self._proc.pid), signal.SIGWINCH)
except OSError:
pass
[docs]
def stop(self):
"""Terminate the subprocess, escalating to kill if needed."""
if self._proc is None:
return
try:
self._proc.terminate()
try:
self._proc.wait(timeout=0.5)
except subprocess.TimeoutExpired:
self._proc.kill()
self._proc.wait(timeout=1.0)
except OSError as e:
log.warning("Error stopping process: %s", e)
self._exit_code = self._proc.returncode
self._cleanup()
[docs]
def write(self, data: str):
"""Write data to the subprocess stdin (or PTY master)."""
raw = data.encode() if isinstance(data, str) else data
if self._master_fd is not None:
try:
os.write(self._master_fd, raw)
except OSError as e:
log.warning("PTY write failed: %s", e)
return
if self._proc is None or self._proc.stdin is None:
return
try:
self._proc.stdin.write(raw)
self._proc.stdin.flush()
except (BrokenPipeError, OSError) as e:
log.warning("ShellNode write failed: %s", e)
[docs]
def on_update(self, dt: float):
if self._proc is None or (reader := self._reader) is None:
return
exited = self._proc.poll() is not None
for name, text in reader.read(exited):
if name == "stdout":
self.stdout_data.emit(text)
else:
self.stderr_data.emit(text)
if not reader.done:
return
self._exit_code = self._proc.returncode
self._cleanup()
self.process_exited.emit(self._exit_code)
def _cleanup(self):
"""Stop reading, close stdin and release the process handle."""
if self._reader is not None:
self._reader.close()
self._reader = None
self._master_fd = None
if self._proc is not None:
if self._proc.stdin:
try:
self._proc.stdin.close()
except OSError:
pass
self._proc = None
[docs]
def on_exit_tree(self):
if self._proc is not None:
self.stop()
[docs]
@property
def running(self) -> bool:
return self._proc is not None and self._proc.poll() is None
[docs]
@property
def exit_code(self) -> int | None:
if self._proc is not None:
return self._proc.returncode
return self._exit_code