Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2patch_stdout
3============
5This implements a context manager that ensures that print statements within
6it won't destroy the user interface. The context manager will replace
7`sys.stdout` by something that draws the output above the current prompt,
8rather than overwriting the UI.
10Usage::
12 with patch_stdout(application):
13 ...
14 application.run()
15 ...
17Multiple applications can run in the body of the context manager, one after the
18other.
19"""
21from __future__ import annotations
23import asyncio
24import queue
25import sys
26import threading
27import time
28from collections.abc import Generator
29from contextlib import contextmanager
30from typing import TextIO, cast
32from .application import get_app_session, run_in_terminal
33from .output import Output
35__all__ = [
36 "patch_stdout",
37 "StdoutProxy",
38]
41@contextmanager
42def patch_stdout(raw: bool = False) -> Generator[None, None, None]:
43 """
44 Replace `sys.stdout` and `sys.stderr` by an :class:`_StdoutProxy` instance.
46 Writing to this proxy will make sure that the text appears above the
47 prompt, and that it doesn't destroy the output from the renderer. If no
48 application is curring, the behavior should be identical to writing to
49 `sys.stdout` directly.
51 Warning: If a new event loop is installed using `asyncio.set_event_loop()`,
52 then make sure that the context manager is applied after the event loop
53 is changed. Printing to stdout will be scheduled in the event loop
54 that's active when the context manager is created.
56 Warning: In order for all text to appear above the prompt `stderr` will also
57 be redirected to the stdout proxy.
59 :param raw: (`bool`) When True, vt100 terminal escape sequences are not
60 removed/escaped.
61 """
62 with StdoutProxy(raw=raw) as proxy:
63 original_stdout = sys.stdout
64 original_stderr = sys.stderr
66 # Enter.
67 sys.stdout = cast(TextIO, proxy)
68 sys.stderr = cast(TextIO, proxy)
70 try:
71 yield
72 finally:
73 sys.stdout = original_stdout
74 sys.stderr = original_stderr
77class _Done:
78 "Sentinel value for stopping the stdout proxy."
81class StdoutProxy:
82 """
83 File-like object, which prints everything written to it, output above the
84 current application/prompt. This class is compatible with other file
85 objects and can be used as a drop-in replacement for `sys.stdout` or can
86 for instance be passed to `logging.StreamHandler`.
88 The current application, above which we print, is determined by looking
89 what application currently runs in the `AppSession` that is active during
90 the creation of this instance.
92 This class can be used as a context manager.
94 In order to avoid having to repaint the prompt continuously for every
95 little write, a short delay of `sleep_between_writes` seconds will be added
96 between writes in order to bundle many smaller writes in a short timespan.
97 """
99 def __init__(
100 self,
101 sleep_between_writes: float = 0.2,
102 raw: bool = False,
103 ) -> None:
104 self.sleep_between_writes = sleep_between_writes
105 self.raw = raw
107 self._lock = threading.RLock()
108 self._buffer: list[str] = []
110 # Keep track of the curret app session.
111 self.app_session = get_app_session()
113 # See what output is active *right now*. We should do it at this point,
114 # before this `StdoutProxy` instance is possibly assigned to `sys.stdout`.
115 # Otherwise, if `patch_stdout` is used, and no `Output` instance has
116 # been created, then the default output creation code will see this
117 # proxy object as `sys.stdout`, and get in a recursive loop trying to
118 # access `StdoutProxy.isatty()` which will again retrieve the output.
119 self._output: Output = self.app_session.output
121 # Flush thread
122 self._flush_queue: queue.Queue[str | _Done] = queue.Queue()
123 self._flush_thread = self._start_write_thread()
124 self.closed = False
126 def __enter__(self) -> StdoutProxy:
127 return self
129 def __exit__(self, *args: object) -> None:
130 self.close()
132 def close(self) -> None:
133 """
134 Stop `StdoutProxy` proxy.
136 This will terminate the write thread, make sure everything is flushed
137 and wait for the write thread to finish.
138 """
139 if not self.closed:
140 self._flush_queue.put(_Done())
141 self._flush_thread.join()
142 self.closed = True
144 def _start_write_thread(self) -> threading.Thread:
145 thread = threading.Thread(
146 target=self._write_thread,
147 name="patch-stdout-flush-thread",
148 daemon=True,
149 )
150 thread.start()
151 return thread
153 def _write_thread(self) -> None:
154 done = False
156 while not done:
157 item = self._flush_queue.get()
159 if isinstance(item, _Done):
160 break
162 # Don't bother calling when we got an empty string.
163 if not item:
164 continue
166 text = []
167 text.append(item)
169 # Read the rest of the queue if more data was queued up.
170 while True:
171 try:
172 item = self._flush_queue.get_nowait()
173 except queue.Empty:
174 break
175 else:
176 if isinstance(item, _Done):
177 done = True
178 else:
179 text.append(item)
181 app_loop = self._get_app_loop()
182 self._write_and_flush(app_loop, "".join(text))
184 # If an application was running that requires repainting, then wait
185 # for a very short time, in order to bundle actual writes and avoid
186 # having to repaint to often.
187 if app_loop is not None:
188 time.sleep(self.sleep_between_writes)
190 def _get_app_loop(self) -> asyncio.AbstractEventLoop | None:
191 """
192 Return the event loop for the application currently running in our
193 `AppSession`.
194 """
195 app = self.app_session.app
197 if app is None:
198 return None
200 return app.loop
202 def _write_and_flush(
203 self, loop: asyncio.AbstractEventLoop | None, text: str
204 ) -> None:
205 """
206 Write the given text to stdout and flush.
207 If an application is running, use `run_in_terminal`.
208 """
210 def write_and_flush() -> None:
211 # Ensure that autowrap is enabled before calling `write`.
212 # XXX: On Windows, the `Windows10_Output` enables/disables VT
213 # terminal processing for every flush. It turns out that this
214 # causes autowrap to be reset (disabled) after each flush. So,
215 # we have to enable it again before writing text.
216 self._output.enable_autowrap()
218 if self.raw:
219 self._output.write_raw(text)
220 else:
221 self._output.write(text)
223 self._output.flush()
225 def write_and_flush_in_loop() -> None:
226 # If an application is running, use `run_in_terminal`, otherwise
227 # call it directly.
228 run_in_terminal(write_and_flush, in_executor=False)
230 if loop is None:
231 # No loop, write immediately.
232 write_and_flush()
233 else:
234 # Make sure `write_and_flush` is executed *in* the event loop, not
235 # in another thread.
236 loop.call_soon_threadsafe(write_and_flush_in_loop)
238 def _write(self, data: str) -> None:
239 """
240 Note: print()-statements cause to multiple write calls.
241 (write('line') and write('\n')). Of course we don't want to call
242 `run_in_terminal` for every individual call, because that's too
243 expensive, and as long as the newline hasn't been written, the
244 text itself is again overwritten by the rendering of the input
245 command line. Therefor, we have a little buffer which holds the
246 text until a newline is written to stdout.
247 """
248 if "\n" in data:
249 # When there is a newline in the data, write everything before the
250 # newline, including the newline itself.
251 before, after = data.rsplit("\n", 1)
252 to_write = self._buffer + [before, "\n"]
253 self._buffer = [after]
255 text = "".join(to_write)
256 self._flush_queue.put(text)
257 else:
258 # Otherwise, cache in buffer.
259 self._buffer.append(data)
261 def _flush(self) -> None:
262 text = "".join(self._buffer)
263 self._buffer = []
264 self._flush_queue.put(text)
266 def write(self, data: str) -> int:
267 with self._lock:
268 self._write(data)
270 return len(data) # Pretend everything was written.
272 def flush(self) -> None:
273 """
274 Flush buffered output.
275 """
276 with self._lock:
277 self._flush()
279 @property
280 def original_stdout(self) -> TextIO | None:
281 return self._output.stdout or sys.__stdout__
283 # Attributes for compatibility with sys.__stdout__:
285 def fileno(self) -> int:
286 return self._output.fileno()
288 def isatty(self) -> bool:
289 stdout = self._output.stdout
290 if stdout is None:
291 return False
293 return stdout.isatty()
295 @property
296 def encoding(self) -> str:
297 return self._output.encoding()
299 @property
300 def errors(self) -> str:
301 return "strict"