1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586 |
- import queue
- import threading
- import time
- from concurrent.futures import Executor, Future
-
-
- class _WorkItem(object):
- """
- 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, fn, args, kwargs):
- self.future = future
- self.fn = fn
- self.args = args
- self.kwargs = kwargs
-
- def run(self):
- 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
- 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):
- self._work_thread = threading.current_thread()
- self._work_queue = queue.Queue()
- self._broken = False
-
- def run_until_future(self, future):
- """
- 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"
- )
- # Keep getting work items and checking the future
- try:
- while True:
- # Get a work item and run it
- try:
- work_item = self._work_queue.get(block=False)
- except queue.Empty:
- # See if the future is done (we only exit if the work queue is empty)
- if future.done():
- return
- # Prevent hot-looping on nothing
- time.sleep(0.001)
- else:
- work_item.run()
- del work_item
- finally:
- self._broken = True
-
- def submit(self, fn, *args, **kwargs):
- # 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()
- work_item = _WorkItem(f, fn, args, kwargs)
- self._work_queue.put(work_item)
- # Return the future
- return f
|