Refactoring of StdoutProxy to take app_session from the caller

examples/ssh/asyncssh-server.py

@@ -11,6 +11,7 @@
 from prompt_toolkit.completion import WordCompleter
 from prompt_toolkit.contrib.ssh import PromptToolkitSSHServer, PromptToolkitSSHSession
 from prompt_toolkit.lexers import PygmentsLexer
+from prompt_toolkit.patch_stdout import StdoutProxy
 from prompt_toolkit.shortcuts import ProgressBar, print_formatted_text
 from prompt_toolkit.shortcuts.dialogs import input_dialog, yes_no_dialog
 from prompt_toolkit.shortcuts.prompt import PromptSession
@@ -99,6 +100,29 @@ async def interact(ssh_session: PromptToolkitSSHSession) -> None:
     await prompt_session.prompt_async("Showing input dialog... [ENTER]")
     await input_dialog("Input dialog", "Running over asyncssh").run_async()
 
+    async def print_counter(output):
+        """
+        Coroutine that prints counters.
+        """
+        try:
+            i = 0
+            while True:
+                output.write(f"Counter: {i}\n")
+                i += 1
+                await asyncio.sleep(3)
+        except asyncio.CancelledError:
+            print("Background task cancelled.")
+
+    with StdoutProxy(app_session=prompt_session) as output:
+        background_task = asyncio.create_task(print_counter(output))
+        try:
+            text = await prompt_session.prompt_async(
+                "Type something with background task: "
+            )
+            output.write(f"You typed: {text}\n")
+        finally:
+            background_task.cancel()
+
 
 async def main(port=8222):
     # Set up logging.

src/prompt_toolkit/application/run_in_terminal.py

@@ -20,7 +20,10 @@
 
 
 def run_in_terminal(
-    func: Callable[[], _T], render_cli_done: bool = False, in_executor: bool = False
+    func: Callable[[], _T],
+    render_cli_done: bool = False,
+    in_executor: bool = False,
+    app=None,
 ) -> Awaitable[_T]:
     """
     Run function on the terminal above the current application or prompt.
@@ -40,12 +43,13 @@ def run_in_terminal(
             erase the interface first.
     :param in_executor: When True, run in executor. (Use this for long
         blocking functions, when you don't want to block the event loop.)
+    :param app: instance of Application. (default None)
 
     :returns: A `Future`.
     """
 
     async def run() -> _T:
-        async with in_terminal(render_cli_done=render_cli_done):
+        async with in_terminal(render_cli_done=render_cli_done, app=app):
             if in_executor:
                 return await run_in_executor_with_context(func)
             else:
@@ -55,7 +59,9 @@ async def run() -> _T:
 
 
 @asynccontextmanager
-async def in_terminal(render_cli_done: bool = False) -> AsyncGenerator[None, None]:
+async def in_terminal(
+    render_cli_done: bool = False, app=None
+) -> AsyncGenerator[None, None]:
     """
     Asynchronous context manager that suspends the current application and runs
     the body in the terminal.
@@ -67,7 +73,8 @@ async def f():
                 call_some_function()
                 await call_some_async_function()
     """
-    app = get_app_or_none()
+    if not app:
+        app = get_app_or_none()
     if app is None or not app._is_running:
         yield
         return

src/prompt_toolkit/patch_stdout.py

@@ -25,10 +25,11 @@
 import threading
 import time
 from contextlib import contextmanager
-from typing import Generator, TextIO, cast
+from typing import Any, Generator, TextIO, cast
 
 from .application import get_app_session, run_in_terminal
 from .output import Output
+from .shortcuts.prompt import PromptSession
 
 __all__ = [
     "patch_stdout",
@@ -95,6 +96,7 @@ def __init__(
         self,
         sleep_between_writes: float = 0.2,
         raw: bool = False,
+        app_session: PromptSession[Any] | None = None,
     ) -> None:
         self.sleep_between_writes = sleep_between_writes
         self.raw = raw
@@ -103,7 +105,9 @@ def __init__(
         self._buffer: list[str] = []
 
         # Keep track of the curret app session.
-        self.app_session = get_app_session()
+        self.app_session = app_session
+        if not self.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`.
@@ -220,7 +224,9 @@ def write_and_flush() -> None:
         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)
+            run_in_terminal(
+                write_and_flush, in_executor=False, app=self.app_session.app
+            )
 
         if loop is None:
             # No loop, write immediately.