1 files changed, 101 insertions, 0 deletions
diff --git a/src/prompt_toolkit/eventloop/utils.py b/src/prompt_toolkit/eventloop/utils.py
new file mode 100644
index 0000000..3138361
--- /dev/null
+++ b/src/prompt_toolkit/eventloop/utils.py
@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import asyncio
+import contextvars
+import sys
+import time
+from asyncio import get_running_loop
+from types import TracebackType
+from typing import Any, Awaitable, Callable, TypeVar, cast
+
+__all__ = [
+    "run_in_executor_with_context",
+    "call_soon_threadsafe",
+    "get_traceback_from_context",
+]
+
+_T = TypeVar("_T")
+
+
+def run_in_executor_with_context(
+    func: Callable[..., _T],
+    *args: Any,
+    loop: asyncio.AbstractEventLoop | None = None,
+) -> Awaitable[_T]:
+    """
+    Run a function in an executor, but make sure it uses the same contextvars.
+    This is required so that the function will see the right application.
+
+    See also: https://bugs.python.org/issue34014
+    """
+    loop = loop or get_running_loop()
+    ctx: contextvars.Context = contextvars.copy_context()
+
+    return loop.run_in_executor(None, ctx.run, func, *args)
+
+
+def call_soon_threadsafe(
+    func: Callable[[], None],
+    max_postpone_time: float | None = None,
+    loop: asyncio.AbstractEventLoop | None = None,
+) -> None:
+    """
+    Wrapper around asyncio's `call_soon_threadsafe`.
+
+    This takes a `max_postpone_time` which can be used to tune the urgency of
+    the method.
+
+    Asyncio runs tasks in first-in-first-out. However, this is not what we
+    want for the render function of the prompt_toolkit UI. Rendering is
+    expensive, but since the UI is invalidated very often, in some situations
+    we render the UI too often, so much that the rendering CPU usage slows down
+    the rest of the processing of the application.  (Pymux is an example where
+    we have to balance the CPU time spend on rendering the UI, and parsing
+    process output.)
+    However, we want to set a deadline value, for when the rendering should
+    happen. (The UI should stay responsive).
+    """
+    loop2 = loop or get_running_loop()
+
+    # If no `max_postpone_time` has been given, schedule right now.
+    if max_postpone_time is None:
+        loop2.call_soon_threadsafe(func)
+        return
+
+    max_postpone_until = time.time() + max_postpone_time
+
+    def schedule() -> None:
+        # When there are no other tasks scheduled in the event loop. Run it
+        # now.
+        # Notice: uvloop doesn't have this _ready attribute. In that case,
+        #         always call immediately.
+        if not getattr(loop2, "_ready", []):
+            func()
+            return
+
+        # If the timeout expired, run this now.
+        if time.time() > max_postpone_until:
+            func()
+            return
+
+        # Schedule again for later.
+        loop2.call_soon_threadsafe(schedule)
+
+    loop2.call_soon_threadsafe(schedule)
+
+
+def get_traceback_from_context(context: dict[str, Any]) -> TracebackType | None:
+    """
+    Get the traceback object from the context.
+    """
+    exception = context.get("exception")
+    if exception:
+        if hasattr(exception, "__traceback__"):
+            return cast(TracebackType, exception.__traceback__)
+        else:
+            # call_exception_handler() is usually called indirectly
+            # from an except block. If it's not the case, the traceback
+            # is undefined...
+            return sys.exc_info()[2]
+
+    return None