Skip to content

Commit 49f9667

Browse files
authored
gh-135736: Fix interaction between TaskGroup and aclose() (#154649)
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 11d0da5 commit 49f9667

4 files changed

Lines changed: 94 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.15
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 exc
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: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1227,6 +1227,72 @@ 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) as cm:
1254+
await fn()
1255+
self.assertEqual(get_error_types(cm.exception), {GeneratorExit})
1256+
1257+
async def test_taskgroup_generator_exit_03(self):
1258+
# A GeneratorExit in one task and an error in another should
1259+
# still give an ExceptionGroup
1260+
async def t1():
1261+
raise GeneratorExit
1262+
1263+
async def t2():
1264+
raise AssertionError('t2 failed')
1265+
1266+
async def fn():
1267+
async with asyncio.TaskGroup() as tg:
1268+
tg.create_task(t1())
1269+
tg.create_task(t2())
1270+
1271+
with self.assertRaises(BaseExceptionGroup) as cm:
1272+
await fn()
1273+
1274+
self.assertEqual(get_error_types(cm.exception), {GeneratorExit, AssertionError})
1275+
1276+
async def test_taskgroup_generator_exit_04(self):
1277+
event = asyncio.Event()
1278+
async def t():
1279+
event.set()
1280+
raise AssertionError('t failed')
1281+
1282+
async def fn():
1283+
async with asyncio.TaskGroup() as tg:
1284+
tg.create_task(t())
1285+
yield 1
1286+
1287+
g = fn()
1288+
await g.asend(None)
1289+
await event.wait() # wait for t() to run
1290+
1291+
with self.assertRaises(BaseExceptionGroup) as cm:
1292+
await g.aclose()
1293+
1294+
self.assertEqual(get_error_types(cm.exception), {GeneratorExit, AssertionError})
1295+
12301296

12311297
class TestTaskGroup(BaseTestTaskGroup, unittest.IsolatedAsyncioTestCase):
12321298
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)