Files
textual-web/src/textual_web/app_session.py

311 lines
9.2 KiB
Python

from __future__ import annotations
import asyncio
from asyncio import StreamReader, StreamWriter, IncompleteReadError
from asyncio.subprocess import Process
from enum import Enum, auto
import io
import logging
import json
import logging
import os
import signal
from time import monotonic
from datetime import timedelta
from pathlib import Path
from importlib_metadata import version
import rich.repr
from . import constants
from .session import Session, SessionConnector
from .types import Meta, SessionID
log = logging.getLogger("textual-web")
class ProcessState(Enum):
"""The state of a process."""
PENDING = auto()
RUNNING = auto()
CLOSING = auto()
CLOSED = auto()
def __repr__(self) -> str:
return self.name
@rich.repr.auto(angular=True)
class AppSession(Session):
"""Runs a single app process."""
def __init__(
self,
working_directory: Path,
command: str,
session_id: SessionID,
devtools: bool = False,
) -> None:
self.working_directory = working_directory
self.command = command
self.session_id = session_id
self.devtools = devtools
self.start_time: float | None = None
self.end_time: float | None = None
self._process: Process | None = None
self._task: asyncio.Task | None = None
super().__init__()
self._state = ProcessState.PENDING
@property
def process(self) -> Process:
"""The asyncio (sub)process"""
assert self._process is not None
return self._process
@property
def stdin(self) -> StreamWriter:
"""The processes stdin."""
assert self._process is not None
assert self._process.stdin is not None
return self._process.stdin
@property
def stdout(self) -> StreamReader:
"""The process' stdout."""
assert self._process is not None
assert self._process.stdout is not None
return self._process.stdout
@property
def stderr(self) -> StreamReader:
"""The process' stderr."""
assert self._process is not None
assert self._process.stderr is not None
return self._process.stderr
@property
def task(self) -> asyncio.Task:
"""Session task."""
assert self._task is not None
return self._task
@property
def state(self) -> ProcessState:
"""Current running state."""
return self._state
@state.setter
def state(self, state: ProcessState) -> None:
self._state = state
run_time = self.run_time
log.debug(
"%r state=%r run_time=%s",
self,
self.state,
"0" if run_time is None else timedelta(seconds=int(run_time)),
)
@property
def run_time(self) -> float | None:
"""Time process was running, or `None` if it hasn't started."""
if self.end_time is not None:
assert self.start_time is not None
return self.end_time - self.start_time
elif self.start_time is not None:
return monotonic() - self.start_time
else:
return None
def __rich_repr__(self) -> rich.repr.Result:
yield self.command
yield "id", self.session_id
if self._process is not None:
yield "returncode", self._process.returncode, None
async def open(self, width: int = 80, height: int = 24) -> None:
"""Open the process."""
environment = dict(os.environ.copy())
environment["TEXTUAL_DRIVER"] = "textual.drivers.web_driver:WebDriver"
environment["TEXTUAL_FPS"] = "60"
environment["TEXTUAL_COLOR_SYSTEM"] = "truecolor"
environment["TERM_PROGRAM"] = "textual-web"
environment["TERM_PROGRAM_VERSION"] = version("textual-web")
environment["COLUMNS"] = str(width)
environment["ROWS"] = str(height)
if self.devtools:
environment["TEXTUAL"] = "debug,devtools"
environment["TEXTUAL_LOG"] = "textual.log"
cwd = os.getcwd()
os.chdir(str(self.working_directory))
try:
self._process = await asyncio.create_subprocess_shell(
self.command,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=environment,
)
finally:
os.chdir(cwd)
await self.set_terminal_size(width, height)
log.debug("opened %r; %r", self.command, self._process)
self.start_time = monotonic()
async def start(self, connector: SessionConnector) -> asyncio.Task:
"""Start a task to run the process."""
self._connector = connector
assert self._task is None
self._task = asyncio.create_task(self.run())
return self._task
async def close(self) -> None:
"""Close the process."""
self.state = ProcessState.CLOSING
await self.send_meta({"type": "quit"})
async def wait(self) -> None:
"""Wait for the process to finish (call close first)."""
if self._task:
await self._task
self._task = None
async def set_terminal_size(self, width: int, height: int) -> None:
"""Set the terminal size for the process.
Args:
width: Width in cells.
height: Height in cells.
"""
await self.send_meta(
{
"type": "resize",
"width": width,
"height": height,
}
)
async def run(self) -> None:
"""This loop reads stdout from the process and relays it through the websocket."""
self.state = ProcessState.RUNNING
META = b"M"
DATA = b"D"
stderr_data = io.BytesIO()
async def read_stderr() -> None:
"""Task to read stderr."""
try:
while True:
data = await self.stderr.read(1024 * 4)
if not data:
break
stderr_data.write(data)
except asyncio.CancelledError:
pass
stderr_task = asyncio.create_task(read_stderr())
readexactly = self.stdout.readexactly
from_bytes = int.from_bytes
on_data = self._connector.on_data
on_meta = self._connector.on_meta
try:
ready = False
for _ in range(10):
line = await self.stdout.readline()
if not line:
break
if line == b"__GANGLION__\n":
ready = True
break
if ready:
while True:
type_bytes = await readexactly(1)
size_bytes = await readexactly(4)
size = from_bytes(size_bytes, "big")
payload = await readexactly(size)
if type_bytes == DATA:
await on_data(payload)
elif type_bytes == META:
meta_data = json.loads(payload)
meta_type = meta_data.get("type")
if meta_type in {"exit", "blur", "focus"}:
await self.send_meta({"type": meta_type})
else:
await on_meta(json.loads(payload))
except IncompleteReadError:
# Incomplete read means that the stream was closed
pass
except asyncio.CancelledError:
pass
finally:
stderr_task.cancel()
await stderr_task
self.end_time = monotonic()
self.state = ProcessState.CLOSED
stderr_message = stderr_data.getvalue().decode("utf-8", errors="replace")
if self._process is not None and self._process.returncode != 0:
if constants.DEBUG and stderr_message:
log.warning(stderr_message)
await self._connector.on_close()
@classmethod
def encode_packet(cls, packet_type: bytes, payload: bytes) -> bytes:
"""Encode a packet.
Args:
packet_type: The packet type (b"D" for data or b"M" for meta)
payload: The payload.
Returns:
Data as bytes.
"""
return b"%s%s%s" % (packet_type, len(payload).to_bytes(4, "big"), payload)
async def send_bytes(self, data: bytes) -> bool:
"""Send bytes to process.
Args:
data: Data to send.
Returns:
True if the data was sent, otherwise False.
"""
stdin = self.stdin
try:
stdin.write(self.encode_packet(b"D", data))
except RuntimeError:
return False
await stdin.drain()
return True
async def send_meta(self, data: Meta) -> bool:
"""Send meta information to process.
Args:
data: Meta dict to send.
Returns:
True if the data was sent, otherwise False.
"""
stdin = self.stdin
data_bytes = json.dumps(data).encode("utf-8")
try:
stdin.write(self.encode_packet(b"M", data_bytes))
except RuntimeError:
return False
await stdin.drain()
return True