1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
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
|
