← Back
Editing: patch_stdout.py
""" patch_stdout ============ This implements a context manager that ensures that print statements within it won't destroy the user interface. The context manager will replace `sys.stdout` by something that draws the output above the current prompt, rather than overwriting the UI. Usage:: with patch_stdout(application): ... application.run() ... Multiple applications can run in the body of the context manager, one after the other. """ import asyncio import queue import sys import threading import time from contextlib import contextmanager from typing import Generator, List, Optional, TextIO, Union, cast from .application import get_app_session, run_in_terminal from .output import Output __all__ = [ "patch_stdout", "StdoutProxy", ] @contextmanager def patch_stdout(raw: bool = False) -> Generator[None, None, None]: """ Replace `sys.stdout` by an :class:`_StdoutProxy` instance. Writing to this proxy will make sure that the text appears above the prompt, and that it doesn't destroy the output from the renderer. If no application is curring, the behaviour should be identical to writing to `sys.stdout` directly. Warning: If a new event loop is installed using `asyncio.set_event_loop()`, then make sure that the context manager is applied after the event loop is changed. Printing to stdout will be scheduled in the event loop that's active when the context manager is created. :param raw: (`bool`) When True, vt100 terminal escape sequences are not removed/escaped. """ with StdoutProxy(raw=raw) as proxy: original_stdout = sys.stdout original_stderr = sys.stderr # Enter. sys.stdout = cast(TextIO, proxy) sys.stderr = cast(TextIO, proxy) try: yield finally: sys.stdout = original_stdout sys.stderr = original_stderr class _Done: "Sentinel value for stopping the stdout proxy." class StdoutProxy: """ File-like object, which prints everything written to it, output above the current application/prompt. This class is compatible with other file objects and can be used as a drop-in replacement for `sys.stdout` or can for instance be passed to `logging.StreamHandler`. The current application, above which we print, is determined by looking what application currently runs in the `AppSession` that is active during the creation of this instance. This class can be used as a context manager. In order to avoid having to repaint the prompt continuously for every little write, a short delay of `sleep_between_writes` seconds will be added between writes in order to bundle many smaller writes in a short timespan. """ def __init__( self, sleep_between_writes: float = 0.2, raw: bool = False, ) -> None: self.sleep_between_writes = sleep_between_writes self.raw = raw self._lock = threading.RLock() self._buffer: List[str] = [] # Keep track of the curret app session. self.app_session = get_app_session() # See what output is active *right now*. We should do it at this point, # before this `StdoutProxy` instance is possibly assigned to `sys.stdout`. # Otherwise, if `patch_stdout` is used, and no `Output` instance has # been created, then the default output creation code will see this # proxy object as `sys.stdout`, and get in a recursive loop trying to # access `StdoutProxy.isatty()` which will again retrieve the output. self._output: Output = self.app_session.output # Flush thread self._flush_queue: queue.Queue[Union[str, _Done]] = queue.Queue() self._flush_thread = self._start_write_thread() self.closed = False def __enter__(self) -> "StdoutProxy": return self def __exit__(self, *args: object) -> None: self.close() def close(self) -> None: """ Stop `StdoutProxy` proxy. This will terminate the write thread, make sure everything is flushed and wait for the write thread to finish. """ if not self.closed: self._flush_queue.put(_Done()) self._flush_thread.join() self.closed = True def _start_write_thread(self) -> threading.Thread: thread = threading.Thread( target=self._write_thread, name="patch-stdout-flush-thread", daemon=True, ) thread.start() return thread def _write_thread(self) -> None: done = False while not done: item = self._flush_queue.get() if isinstance(item, _Done): break # Don't bother calling when we got an empty string. if not item: continue text = [] text.append(item) # Read the rest of the queue if more data was queued up. while True: try: item = self._flush_queue.get_nowait() except queue.Empty: break else: if isinstance(item, _Done): done = True else: text.append(item) app_loop = self._get_app_loop() self._write_and_flush(app_loop, "".join(text)) # If an application was running that requires repainting, then wait # for a very short time, in order to bundle actual writes and avoid # having to repaint to often. if app_loop is not None: time.sleep(self.sleep_between_writes) def _get_app_loop(self) -> Optional[asyncio.AbstractEventLoop]: """ Return the event loop for the application currently running in our `AppSession`. """ app = self.app_session.app if app is None: return None return app.loop def _write_and_flush( self, loop: Optional[asyncio.AbstractEventLoop], text: str ) -> None: """ Write the given text to stdout and flush. If an application is running, use `run_in_terminal`. """ def write_and_flush() -> None: if self.raw: self._output.write_raw(text) else: self._output.write(text) self._output.flush() def write_and_flush_in_loop() -> None: # If an application is running, use `run_in_terminal`, otherwise # call it directly. run_in_terminal(write_and_flush, in_executor=False) if loop is None: # No loop, write immediately. write_and_flush() else: # Make sure `write_and_flush` is executed *in* the event loop, not # in another thread. loop.call_soon_threadsafe(write_and_flush_in_loop) def _write(self, data: str) -> None: """ Note: print()-statements cause to multiple write calls. (write('line') and write('\n')). Of course we don't want to call `run_in_terminal` for every individual call, because that's too expensive, and as long as the newline hasn't been written, the text itself is again overwritten by the rendering of the input command line. Therefor, we have a little buffer which holds the text until a newline is written to stdout. """ if "\n" in data: # When there is a newline in the data, write everything before the # newline, including the newline itself. before, after = data.rsplit("\n", 1) to_write = self._buffer + [before, "\n"] self._buffer = [after] text = "".join(to_write) self._flush_queue.put(text) else: # Otherwise, cache in buffer. self._buffer.append(data) def _flush(self) -> None: text = "".join(self._buffer) self._buffer = [] self._flush_queue.put(text) def write(self, data: str) -> int: with self._lock: self._write(data) return len(data) # Pretend everything was written. def flush(self) -> None: """ Flush buffered output. """ with self._lock: self._flush() @property def original_stdout(self) -> TextIO: return self._output.stdout or sys.__stdout__ # Attributes for compatibility with sys.__stdout__: def fileno(self) -> int: return self._output.fileno() def isatty(self) -> bool: stdout = self._output.stdout if stdout is None: return False return stdout.isatty() @property def encoding(self) -> str: return self._output.encoding() @property def errors(self) -> str: return "strict"
Save File
Cancel