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