Skip to content

Commit 5b945dd

Browse files
committed
gh-135736: Fix interaction between TaskGroup and aclose()
Currently if you have a generator like: ``` async def test(): async with asyncio.TaskGroup() as tg: async for x in whatever(): yield x ``` If aclose() is called on the generator and GeneratorExit is raised, the TaskGroup's __aexit__ will raise a BaseExceptionGroup, and that will be raised from the aclose() call instead of it being swallowed. Fix this by raising GeneratorExit from __aexit__ when: 1. The body of task group raised it 2. No subtasks raised exceptions This makes GeneratorExit work without swallowing other exceptions (like it would if we just added it to _is_base_error).
1 parent 998fc4a commit 5b945dd

4 files changed

Lines changed: 93 additions & 4 deletions

File tree

Doc/library/asyncio-task.rst

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -433,6 +433,10 @@ unless it is :exc:`asyncio.CancelledError`,
433433
is also included in the exception group.
434434
The same special case is made for
435435
:exc:`KeyboardInterrupt` and :exc:`SystemExit` as in the previous paragraph.
436+
There is an additional special case made only for the body of the
437+
``async with``: if it raises :exc:`GeneratorExit` and none of the
438+
other tasks raise exceptions that would be reported, then the
439+
:exc:`GeneratorExit` is reraised.
436440

437441
Task groups are careful not to mix up the internal cancellation used to
438442
"wake up" their :meth:`~object.__aexit__` with cancellation requests
@@ -456,6 +460,10 @@ reported by :meth:`asyncio.Task.cancelling`.
456460
Improved handling of simultaneous internal and external cancellations
457461
and correct preservation of cancellation counts.
458462

463+
.. versionchanged:: 3.16
464+
465+
Addition of the special case for :exc:`GeneratorExit`.
466+
459467
Sleeping
460468
========
461469

Lib/asyncio/taskgroups.py

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -174,10 +174,23 @@ async def _aexit(self, et, exc):
174174
self._parent_task.uncancel()
175175
self._parent_task.cancel()
176176
try:
177-
raise BaseExceptionGroup(
178-
'unhandled errors in a TaskGroup',
179-
self._errors,
180-
) from None
177+
# If the *only* error is a GeneratorExit from the body
178+
# of the group, then instead of raising an
179+
# ExceptionGroup we raise GeneratorExit. This ensures
180+
# that async generators that use TaskGroup properly
181+
# swallow the exception on `aclose()` while ensuring
182+
# that no exceptions from subtasks are swallowed.
183+
if (
184+
et is not None
185+
and issubclass(et, GeneratorExit)
186+
and len(self._errors) == 1
187+
):
188+
raise GeneratorExit
189+
else:
190+
raise BaseExceptionGroup(
191+
'unhandled errors in a TaskGroup',
192+
self._errors,
193+
) from None
181194
finally:
182195
exc = None
183196

Lib/test/test_asyncio/test_taskgroups.py

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1227,6 +1227,71 @@ async def fn_3():
12271227
self.assertEqual(await race(fn_1, fn_2, fn_3), 1)
12281228
self.assertListEqual(record, ["1 started", "2 started", "3 started", "1 finished"])
12291229

1230+
async def test_taskgroup_generator_exit_01(self):
1231+
# GeneratorExit in a TaskGroup should be fine
1232+
async def gen():
1233+
yield 1
1234+
1235+
async def fn():
1236+
async with asyncio.TaskGroup() as tg:
1237+
async for n in gen():
1238+
yield n
1239+
1240+
g = fn()
1241+
await g.asend(None)
1242+
await g.aclose()
1243+
1244+
async def test_taskgroup_generator_exit_02(self):
1245+
# A lone GeneratorExit in a task should still give an ExceptionGroup
1246+
async def t():
1247+
raise GeneratorExit
1248+
1249+
async def fn():
1250+
async with asyncio.TaskGroup() as tg:
1251+
tg.create_task(t())
1252+
1253+
with self.assertRaises(BaseExceptionGroup):
1254+
await fn()
1255+
1256+
async def test_taskgroup_generator_exit_03(self):
1257+
# A GeneratorExit in one task and an error in another should
1258+
# still give an ExceptionGroup
1259+
async def t1():
1260+
raise GeneratorExit
1261+
1262+
async def t2():
1263+
raise AssertionError('t2 failed')
1264+
1265+
async def fn():
1266+
async with asyncio.TaskGroup() as tg:
1267+
tg.create_task(t1())
1268+
tg.create_task(t2())
1269+
1270+
with self.assertRaises(BaseExceptionGroup) as cm:
1271+
await fn()
1272+
1273+
self.assertEqual(get_error_types(cm.exception), {GeneratorExit, AssertionError})
1274+
1275+
async def test_taskgroup_generator_exit_04(self):
1276+
event = asyncio.Event()
1277+
async def t():
1278+
event.set()
1279+
raise AssertionError('t failed')
1280+
1281+
async def fn():
1282+
async with asyncio.TaskGroup() as tg:
1283+
tg.create_task(t())
1284+
yield 1
1285+
1286+
g = fn()
1287+
await g.asend(None)
1288+
await event.wait() # wait for t() to run
1289+
1290+
with self.assertRaises(BaseExceptionGroup) as cm:
1291+
await g.aclose()
1292+
1293+
self.assertEqual(get_error_types(cm.exception), {GeneratorExit, AssertionError})
1294+
12301295

12311296
class TestTaskGroup(BaseTestTaskGroup, unittest.IsolatedAsyncioTestCase):
12321297
loop_factory = asyncio.EventLoop
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Fix ::class:`asyncio.TaskGroup` to not wrap a :exc:`GeneratorExit` into a
2+
:exc:`BaseExceptionGroup` if it was raised by the body of the task group and
3+
none of the tasks in the group raised exceptions.

0 commit comments

Comments
 (0)