summaryrefslogtreecommitdiffstats
path: root/src/prompt_toolkit/eventloop/utils.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/prompt_toolkit/eventloop/utils.py')
-rw-r--r--src/prompt_toolkit/eventloop/utils.py118
1 files changed, 118 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..2e5a05e
--- /dev/null
+++ b/src/prompt_toolkit/eventloop/utils.py
@@ -0,0 +1,118 @@
+import asyncio
+import sys
+import time
+from types import TracebackType
+from typing import Any, Awaitable, Callable, Dict, Optional, TypeVar, cast
+
+try:
+ import contextvars
+except ImportError:
+ from . import dummy_contextvars as contextvars # type: ignore
+
+__all__ = [
+ "run_in_executor_with_context",
+ "call_soon_threadsafe",
+ "get_traceback_from_context",
+ "get_event_loop",
+]
+
+_T = TypeVar("_T")
+
+
+def run_in_executor_with_context(
+ func: Callable[..., _T],
+ *args: Any,
+ loop: Optional[asyncio.AbstractEventLoop] = 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_event_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: Optional[float] = None,
+ loop: Optional[asyncio.AbstractEventLoop] = 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_event_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]) -> Optional[TracebackType]:
+ """
+ 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
+
+
+def get_event_loop() -> asyncio.AbstractEventLoop:
+ """Backward compatible way to get the event loop"""
+ # Python 3.6 doesn't have get_running_loop
+ # Python 3.10 deprecated get_event_loop
+ if sys.version_info >= (3, 7):
+ getloop = asyncio.get_running_loop
+ else:
+ getloop = asyncio.get_event_loop
+
+ try:
+ return getloop()
+ except RuntimeError:
+ return asyncio.get_event_loop_policy().get_event_loop()