Skip to content

Commit 925dc7f

Browse files
authored
bpo-39606: allow closing async generators that are already closed (GH-18475)
The fix for [bpo-39386](https://bugs.python.org/issue39386) attempted to make it so you couldn't reuse a agen.aclose() coroutine object. It accidentally also prevented you from calling aclose() at all on an async generator that was already closed or exhausted. This commit fixes it so we're only blocking the actually illegal cases, while allowing the legal cases. The new tests failed before this patch. Also confirmed that this fixes the test failures we were seeing in Trio with Python dev builds: python-trio/trio#1396 https://bugs.python.org/issue39606
1 parent 7514f4f commit 925dc7f

File tree

3 files changed

+41
-6
lines changed

3 files changed

+41
-6
lines changed

Lib/test/test_asyncgen.py

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1128,7 +1128,7 @@ async def main():
11281128

11291129
self.assertEqual([], messages)
11301130

1131-
def test_async_gen_await_anext_twice(self):
1131+
def test_async_gen_await_same_anext_coro_twice(self):
11321132
async def async_iterate():
11331133
yield 1
11341134
yield 2
@@ -1147,7 +1147,7 @@ async def run():
11471147

11481148
self.loop.run_until_complete(run())
11491149

1150-
def test_async_gen_await_aclose_twice(self):
1150+
def test_async_gen_await_same_aclose_coro_twice(self):
11511151
async def async_iterate():
11521152
yield 1
11531153
yield 2
@@ -1164,6 +1164,32 @@ async def run():
11641164

11651165
self.loop.run_until_complete(run())
11661166

1167+
def test_async_gen_aclose_twice_with_different_coros(self):
1168+
# Regression test for https://bugs.python.org/issue39606
1169+
async def async_iterate():
1170+
yield 1
1171+
yield 2
1172+
1173+
async def run():
1174+
it = async_iterate()
1175+
await it.aclose()
1176+
await it.aclose()
1177+
1178+
self.loop.run_until_complete(run())
1179+
1180+
def test_async_gen_aclose_after_exhaustion(self):
1181+
# Regression test for https://bugs.python.org/issue39606
1182+
async def async_iterate():
1183+
yield 1
1184+
yield 2
1185+
1186+
async def run():
1187+
it = async_iterate()
1188+
async for _ in it:
1189+
pass
1190+
await it.aclose()
1191+
1192+
self.loop.run_until_complete(run())
11671193

11681194
if __name__ == "__main__":
11691195
unittest.main()
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix regression caused by fix for bpo-39386, that prevented calling
2+
``aclose`` on an async generator that had already been closed or exhausted.

Objects/genobject.c

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1797,16 +1797,22 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
17971797
PyFrameObject *f = gen->gi_frame;
17981798
PyObject *retval;
17991799

1800-
if (f == NULL || f->f_stacktop == NULL ||
1801-
o->agt_state == AWAITABLE_STATE_CLOSED) {
1800+
if (o->agt_state == AWAITABLE_STATE_CLOSED) {
18021801
PyErr_SetString(
18031802
PyExc_RuntimeError,
18041803
"cannot reuse already awaited aclose()/athrow()");
18051804
return NULL;
18061805
}
18071806

1807+
if (f == NULL || f->f_stacktop == NULL) {
1808+
o->agt_state = AWAITABLE_STATE_CLOSED;
1809+
PyErr_SetNone(PyExc_StopIteration);
1810+
return NULL;
1811+
}
1812+
18081813
if (o->agt_state == AWAITABLE_STATE_INIT) {
18091814
if (o->agt_gen->ag_running_async) {
1815+
o->agt_state = AWAITABLE_STATE_CLOSED;
18101816
if (o->agt_args == NULL) {
18111817
PyErr_SetString(
18121818
PyExc_RuntimeError,
@@ -1878,7 +1884,6 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
18781884
/* aclose() mode */
18791885
if (retval) {
18801886
if (_PyAsyncGenWrappedValue_CheckExact(retval)) {
1881-
o->agt_gen->ag_running_async = 0;
18821887
Py_DECREF(retval);
18831888
goto yield_close;
18841889
}
@@ -1893,16 +1898,17 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
18931898

18941899
yield_close:
18951900
o->agt_gen->ag_running_async = 0;
1901+
o->agt_state = AWAITABLE_STATE_CLOSED;
18961902
PyErr_SetString(
18971903
PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG);
18981904
return NULL;
18991905

19001906
check_error:
19011907
o->agt_gen->ag_running_async = 0;
1908+
o->agt_state = AWAITABLE_STATE_CLOSED;
19021909
if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration) ||
19031910
PyErr_ExceptionMatches(PyExc_GeneratorExit))
19041911
{
1905-
o->agt_state = AWAITABLE_STATE_CLOSED;
19061912
if (o->agt_args == NULL) {
19071913
/* when aclose() is called we don't want to propagate
19081914
StopAsyncIteration or GeneratorExit; just raise
@@ -1936,6 +1942,7 @@ async_gen_athrow_throw(PyAsyncGenAThrow *o, PyObject *args)
19361942
/* aclose() mode */
19371943
if (retval && _PyAsyncGenWrappedValue_CheckExact(retval)) {
19381944
o->agt_gen->ag_running_async = 0;
1945+
o->agt_state = AWAITABLE_STATE_CLOSED;
19391946
Py_DECREF(retval);
19401947
PyErr_SetString(PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG);
19411948
return NULL;

0 commit comments

Comments
 (0)