diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 3590748c477a1f..d1b1e8fbb993cd 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -570,28 +570,6 @@ def _wrap_awaitable(awaitable): return (yield from awaitable.__await__()) -class _GatheringFuture(futures.Future): - """Helper for gather(). - - This overrides cancel() to cancel all the children and act more - like Task.cancel(), which doesn't immediately mark itself as - cancelled. - """ - - def __init__(self, children, *, loop=None): - super().__init__(loop=loop) - self._children = children - - def cancel(self): - if self.done(): - return False - ret = False - for child in self._children: - if child.cancel(): - ret = True - return ret - - def gather(*coros_or_futures, loop=None, return_exceptions=False): """Return a future aggregating results from the given coroutines/futures. @@ -615,84 +593,82 @@ def gather(*coros_or_futures, loop=None, return_exceptions=False): prevent the cancellation of one child to cause other children to be cancelled.) """ - if not coros_or_futures: - if loop is None: - loop = events.get_event_loop() - outer = loop.create_future() - outer.set_result([]) - return outer - - def _done_callback(fut): - nonlocal nfinished - nfinished += 1 - - if outer.done(): - if not fut.cancelled(): - # Mark exception retrieved. - fut.exception() - return + for corf in coros_or_futures: + if isinstance(corf, futures.Future): + if loop is None: + loop = corf._loop + elif loop is not corf._loop: + raise ValueError( + "all futures must belong to the same loop in gather") + return ensure_future(_gather(*coros_or_futures, loop=loop, + return_exceptions=return_exceptions, + _nocancel=True), + loop=loop) + + +async def _gather(*coros_or_futures, loop=None, + return_exceptions=False, _nocancel=False): + """A coroutine version of gather + + gather() was traditionally a plain function that returns a future. This is a + coroutine implentation of it, and gather() wraps it in an ensure_future(), + and does some other minor stuff for backwards compatiblity. + + The parameters are the same as for gather(). A new, internal _nocancel + parameter is added for backwards compatiblity: if it is True, the task + returned by gather() will have a CancelledError set if one of the gathered + futures raises it, but the future is not considered cancelled(). + """ + if loop is None: + loop = events.get_event_loop() - if not return_exceptions: - if fut.cancelled(): - # Check if 'fut' is cancelled first, as - # 'fut.exception()' will *raise* a CancelledError - # instead of returning it. - exc = futures.CancelledError() - outer.set_exception(exc) - return - else: - exc = fut.exception() - if exc is not None: - outer.set_exception(exc) - return + async def run_one(awaitable, args): + nonlocal n - if nfinished == nfuts: - # All futures are done; create a list of results - # and set it to the 'outer' future. - results = [] - - for fut in children: - if fut.cancelled(): - # Check if 'fut' is cancelled first, as - # 'fut.exception()' will *raise* a CancelledError - # instead of returning it. - res = futures.CancelledError() - else: - res = fut.exception() - if res is None: - res = fut.result() - results.append(res) - - outer.set_result(results) - - arg_to_fut = {} - children = [] - nfuts = 0 - nfinished = 0 - for arg in coros_or_futures: - if arg not in arg_to_fut: - fut = ensure_future(arg, loop=loop) - if loop is None: - loop = futures._get_loop(fut) - if fut is not arg: - # 'arg' was not a Future, therefore, 'fut' is a new - # Future created specifically for 'arg'. Since the caller - # can't control it, disable the "destroy pending task" - # warning. - fut._log_destroy_pending = False - - nfuts += 1 - arg_to_fut[arg] = fut - fut.add_done_callback(_done_callback) + try: + ret = await awaitable + for i in args: + results[i] = ret + except Exception as exc: + if return_exceptions: + for i in args: + results[i] = exc + elif not future.done(): + future.set_exception(exc) + finally: + n -= 1 + if n == 0 and not future.done(): + future.set_result(None) - else: - # There's a duplicate Future object in coros_or_futures. - fut = arg_to_fut[arg] + unique = {} + for i, arg in enumerate(coros_or_futures): + unique.setdefault(arg, []).append(i) - children.append(fut) + tasks = [ensure_future(run_one(k, v), loop=loop) + for k, v in unique.items()] + n = len(tasks) + results = [None] * len(coros_or_futures) - outer = _GatheringFuture(children, loop=loop) - return outer + while n > 0: + future = loop.create_future() + try: + await future + except futures.CancelledError as exc: + if future.cancelled(): + if n == 0: + # Python issue #26923: asyncio.gather drops cancellation + raise + for task in tasks: + task.cancel() + else: # one of the gathered futures got cancelled + if _nocancel: + # this sets a CancelledError for the task, but + # task.cancelled() stays False. + current_task(loop=loop).cancel() + return + else: + raise + return results def shield(arg, *, loop=None): diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index e5334c6ff914c6..605ca05947ee6d 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -1714,8 +1714,8 @@ def test_gather_shield(self): test_utils.run_briefly(self.loop) parent.cancel() # This should cancel inner1 and inner2 but bot child1 and child2. - test_utils.run_briefly(self.loop) - self.assertIsInstance(parent.exception(), asyncio.CancelledError) + with self.assertRaises(asyncio.CancelledError): + self.loop.run_until_complete(parent) self.assertTrue(inner1.cancelled()) self.assertTrue(inner2.cancelled()) child1.set_result(1) @@ -2036,34 +2036,26 @@ def test_cancel_wait_for(self): self._test_cancel_wait_for(60.0) def test_cancel_gather(self): - """Ensure that a gathering future refuses to be cancelled once all - children are done""" + """test gather cancels propelry after children are done + + This is Python issue #26923 + """ loop = asyncio.new_event_loop() self.addCleanup(loop.close) fut = self.new_future(loop) - # The indirection fut->child_coro is needed since otherwise the - # gathering task is done at the same time as the child future - def child_coro(): - return (yield from fut) - gather_future = asyncio.gather(child_coro(), loop=loop) - gather_task = asyncio.ensure_future(gather_future, loop=loop) - - cancel_result = None - def cancelling_callback(_): - nonlocal cancel_result - cancel_result = gather_task.cancel() - fut.add_done_callback(cancelling_callback) - - fut.set_result(42) # calls the cancelling_callback after fut is done() - - # At this point the task should complete. - loop.run_until_complete(gather_task) - - # Python issue #26923: asyncio.gather drops cancellation - self.assertEqual(cancel_result, False) - self.assertFalse(gather_task.cancelled()) - self.assertEqual(gather_task.result(), [42]) + + async def killer(): + fut.set_result(42) + task.cancel() + + async def outer(): + await asyncio.gather(fut, killer(), loop=loop) + self.fail('cancelled gather did not cancel its task') + + task = asyncio.ensure_future(outer(), loop=loop) + with self.assertRaises(asyncio.CancelledError): + loop.run_until_complete(task) def test_exception_traceback(self): # See http://bugs.python.org/issue28843