def wait_for(fut: FutureOrCoroutine[T], timeout: float) -> AbstractFuture[T]:
"""
Returns a future that will be cancelled after timeout seconds have
passed.
If timeout is zero, it will not timeout.
:param fut: The future to wait for.
:param timeout: When should the execution be cancelled.
:return: A new future that will be cancelled after the given timeout.
"""
def _expire(_):
"""Cancel future on expiry; fire a timeout exception."""
if not fut.done():
fut.cancel()
result.set_exception(TimeoutError())
def _complete(_):
"""Cancel the timeouter when future within the timeout."""
timeouter.cancel()
result: AbstractFuture[T] = Future()
timeouter = sleep(timeout) if timeout else Future()
fut = wrap_future(fut)
timeouter.add_done_callback(_expire)
result.add_done_callback(_complete)
copy(fut, result, copy_cancel=False)
return result
def test_no_result_copy_cancel(self):
source = Future()
target = Future()
copy(source, target, copy_result=False)
source.cancel()
self.assertTrue(target.cancel())
self.assertTrue(target.done())
def test_no_result_copy_exception(self):
source = Future()
target = Future()
copy(source, target, copy_result=False)
source.set_exception(exc)
self.assertFalse(target.cancelled())
self.assertFalse(target.done())
def test_no_result_copy_success(self):
source = Future()
target = Future()
copy(source, target, copy_result=False)
source.set_result(obj)
self.assertFalse(target.cancelled())
self.assertFalse(target.done())
def test_copy_cancel(self):
source = Future()
target = Future()
copy(source, target)
source.cancel()
self.assertTrue(target.cancel())
self.assertTrue(target.done())
def test_no_cancel_copy_cancel(self):
source = Future()
target = Future()
copy(source, target, copy_cancel=False)
source.cancel()
self.assertFalse(target.cancelled())
self.assertFalse(target.done())
def gather(*futs_or_coros: FutureOrCoroutine[T],
return_exceptions: bool = False) -> AbstractFuture[Sequence[T]]:
"""
Gathers the results of the exceptions.
:param futs_or_coros: The futures whose results are to be gathered.
:param return_exceptions: If false, if a future rejects, the gather future will reject.
:return: A future that gathers the results.
"""
if return_exceptions:
wait_mode = ALL_COMPLETED
else:
wait_mode = FIRST_EXCEPTION
futs = list(map(wrap_future, futs_or_coros))
def _propagate_cancel(_):
if not result.cancelled():
return
for fut in futs:
if fut.done():
continue
fut.cancel()
def _gather_result(_):
result_list = []
for fut in futs:
if fut.cancelled():
if return_exceptions:
result_list.append(CancelledError())
continue
result.set_exception(CancelledError())
break
if fut.exception() is None:
result_list.append(fut.result())
continue
if return_exceptions:
result_list.append(fut.exception())
continue
result.set_exception(fut.exception())
break
result.set_result(result_list)
result: AbstractFuture[Sequence[T]] = Future()
result.add_done_callback(_propagate_cancel)
waiter: AbstractFuture[DoneAndNotDoneFutures] = wait(futs,
return_when=wait_mode)
waiter.add_done_callback(_gather_result)
copy(result, waiter, copy_result=False)
return result
def test_copy_exception(self):
source = Future()
target = Future()
copy(source, target)
source.set_exception(exc)
self.assertFalse(target.cancelled())
self.assertTrue(target.done())
self.assertIs(target.exception(), exc)
def test_copy_success(self):
source = Future()
target = Future()
copy(source, target)
source.set_result(obj)
self.assertFalse(target.cancelled())
self.assertTrue(target.done())
self.assertIsNone(target.exception())
self.assertIs(target.result(), obj)
def shield(fut: FutureOrCoroutine[T]) -> AbstractFuture[T]:
"""
Shields a future from being cancelled by the parent.
If the future is cancelled that is passed to this function, the
future returned by this function will reject with a
:class:`concurrent.futures.CancelledError`.
:param fut: The future to shield.
:return: A future that will not propagate its cancellation to the parent.
"""
def _bubble_child(_):
if target.cancelled():
pass
if fut.cancelled():
target.set_exception(CancelledError())
target: AbstractFuture[T] = Future()
copy(fut, wrap_future(target), copy_cancel=False)
fut.add_done_callback(_bubble_child)
return target
def wait(futs_or_coros: Sequence[FutureOrCoroutine[T]],
timeout: Real = 0,
return_when=ALL_COMPLETED) -> AbstractFuture[DoneAndNotDoneFutures]:
"""
Waits for multiple futures.
If `return_when` is ALL_COMPLETED it will wait until
all passed futures have finished.
If `return_when` is FIRST_COMPLETED it will wait until
the first future resolves or rejects.
If `return_when` is FIRST_EXCEPTION it will wait until
the first future rejects and otherwise wait until all
futures finished.
If `timeout` is greater than 0, it will wait maximally
for the given amount of time. Then it will resolve.
In all cases a named-tuple with `done` containing all futures
that have finished until the wait-future resolved and `not_done`
with all futures that still run.
:param futs_or_coros: The futures and/or coroutines to wait for.
:param timeout: The maximal time to wait for them.
:param return_when: When to return.
:return: A future that will resolve-conditions have passed.
"""
cond = Future()
result = Future()
finished = []
running = list(map(wrap_future, futs_or_coros))
lock = Lock()
def _single_finishes(fut: AbstractFuture[T]):
if cond.done():
return
with lock:
running.remove(fut)
if fut.cancelled():
fut = reject(CancelledError())
finished.append(fut)
empty = not running
if return_when == FIRST_COMPLETED:
cond.set_result(None)
return
if fut.exception() and return_when == FIRST_EXCEPTION:
cond.set_result(None)
return
if empty:
cond.set_result(None)
for f in running[:]:
f.add_done_callback(_single_finishes)
def _timeout(_):
if result.done():
return
cond.set_result(None)
def _completes(_):
if result.cancelled():
return
result.set_result(
DoneAndNotDoneFutures(done=finished, not_done=running))
timeouter = sleep(timeout, None) if timeout else Future()
timeouter.add_done_callback(_timeout)
copy(result, timeouter, copy_result=False)
# Copy cancel state to cond.
copy(result, cond, copy_result=False)
cond.add_done_callback(_completes)
return result