diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-05-04 17:35:20 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-05-04 17:35:20 +0000 |
commit | e106bf94eff07d9a59771d9ccc4406421e18ab64 (patch) | |
tree | edb6545500e39df9c67aa918a6125bffc8ec1aee /src/prompt_toolkit/eventloop/utils.py | |
parent | Initial commit. (diff) | |
download | prompt-toolkit-upstream.tar.xz prompt-toolkit-upstream.zip |
Adding upstream version 3.0.36.upstream/3.0.36upstream
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to '')
-rw-r--r-- | src/prompt_toolkit/eventloop/utils.py | 118 |
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() |