diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 15422da1b3..68eafc96bf 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -816,7 +816,7 @@ def shield(arg, *, loop=None): loop = futures._get_loop(inner) outer = loop.create_future() - def _done_callback(inner): + def _inner_done_callback(inner): if outer.cancelled(): if not inner.cancelled(): # Mark inner's result as retrieved. @@ -832,7 +832,13 @@ def shield(arg, *, loop=None): else: outer.set_result(inner.result()) - inner.add_done_callback(_done_callback) + + def _outer_done_callback(outer): + if not inner.done(): + inner.remove_done_callback(_inner_done_callback) + + inner.add_done_callback(_inner_done_callback) + outer.add_done_callback(_outer_done_callback) return outer diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index c65d1f2440..458ef3ff65 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -1768,7 +1768,7 @@ class BaseTaskTests: test_utils.run_briefly(self.loop) self.assertIs(outer.exception(), exc) - def test_shield_cancel(self): + def test_shield_cancel_inner(self): inner = self.new_future(self.loop) outer = asyncio.shield(inner) test_utils.run_briefly(self.loop) @@ -1776,6 +1776,15 @@ class BaseTaskTests: test_utils.run_briefly(self.loop) self.assertTrue(outer.cancelled()) + def test_shield_cancel_outer(self): + inner = self.new_future(self.loop) + outer = asyncio.shield(inner) + test_utils.run_briefly(self.loop) + outer.cancel() + test_utils.run_briefly(self.loop) + self.assertTrue(outer.cancelled()) + self.assertEqual(0, 0 if outer._callbacks is None else len(outer._callbacks)) + def test_shield_shortcut(self): fut = self.new_future(self.loop) fut.set_result(42)