116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
import queue
|
|
import sys
|
|
import threading
|
|
from concurrent.futures import Executor, Future
|
|
from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union
|
|
|
|
if sys.version_info >= (3, 10):
|
|
from typing import ParamSpec
|
|
else:
|
|
from typing_extensions import ParamSpec
|
|
|
|
_T = TypeVar("_T")
|
|
_P = ParamSpec("_P")
|
|
_R = TypeVar("_R")
|
|
|
|
|
|
class _WorkItem:
|
|
"""
|
|
Represents an item needing to be run in the executor.
|
|
Copied from ThreadPoolExecutor (but it's private, so we're not going to rely on importing it)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
future: "Future[_R]",
|
|
fn: Callable[_P, _R],
|
|
*args: _P.args,
|
|
**kwargs: _P.kwargs,
|
|
):
|
|
self.future = future
|
|
self.fn = fn
|
|
self.args = args
|
|
self.kwargs = kwargs
|
|
|
|
def run(self) -> None:
|
|
__traceback_hide__ = True # noqa: F841
|
|
if not self.future.set_running_or_notify_cancel():
|
|
return
|
|
try:
|
|
result = self.fn(*self.args, **self.kwargs)
|
|
except BaseException as exc:
|
|
self.future.set_exception(exc)
|
|
# Break a reference cycle with the exception 'exc'
|
|
self = None # type: ignore[assignment]
|
|
else:
|
|
self.future.set_result(result)
|
|
|
|
|
|
class CurrentThreadExecutor(Executor):
|
|
"""
|
|
An Executor that actually runs code in the thread it is instantiated in.
|
|
Passed to other threads running async code, so they can run sync code in
|
|
the thread they came from.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._work_thread = threading.current_thread()
|
|
self._work_queue: queue.Queue[Union[_WorkItem, "Future[Any]"]] = queue.Queue()
|
|
self._broken = False
|
|
|
|
def run_until_future(self, future: "Future[Any]") -> None:
|
|
"""
|
|
Runs the code in the work queue until a result is available from the future.
|
|
Should be run from the thread the executor is initialised in.
|
|
"""
|
|
# Check we're in the right thread
|
|
if threading.current_thread() != self._work_thread:
|
|
raise RuntimeError(
|
|
"You cannot run CurrentThreadExecutor from a different thread"
|
|
)
|
|
future.add_done_callback(self._work_queue.put)
|
|
# Keep getting and running work items until we get the future we're waiting for
|
|
# back via the future's done callback.
|
|
try:
|
|
while True:
|
|
# Get a work item and run it
|
|
work_item = self._work_queue.get()
|
|
if work_item is future:
|
|
return
|
|
assert isinstance(work_item, _WorkItem)
|
|
work_item.run()
|
|
del work_item
|
|
finally:
|
|
self._broken = True
|
|
|
|
def _submit(
|
|
self,
|
|
fn: Callable[_P, _R],
|
|
*args: _P.args,
|
|
**kwargs: _P.kwargs,
|
|
) -> "Future[_R]":
|
|
# Check they're not submitting from the same thread
|
|
if threading.current_thread() == self._work_thread:
|
|
raise RuntimeError(
|
|
"You cannot submit onto CurrentThreadExecutor from its own thread"
|
|
)
|
|
# Check they're not too late or the executor errored
|
|
if self._broken:
|
|
raise RuntimeError("CurrentThreadExecutor already quit or is broken")
|
|
# Add to work queue
|
|
f: "Future[_R]" = Future()
|
|
work_item = _WorkItem(f, fn, *args, **kwargs)
|
|
self._work_queue.put(work_item)
|
|
# Return the future
|
|
return f
|
|
|
|
# Python 3.9+ has a new signature for submit with a "/" after `fn`, to enforce
|
|
# it to be a positional argument. If we ignore[override] mypy on 3.9+ will be
|
|
# happy but 3.7/3.8 will say that the ignore comment is unused, even when
|
|
# defining them differently based on sys.version_info.
|
|
# We should be able to remove this when we drop support for 3.7/3.8.
|
|
if not TYPE_CHECKING:
|
|
|
|
def submit(self, fn, *args, **kwargs):
|
|
return self._submit(fn, *args, **kwargs)
|