Skip to content

Commit fc1732c

Browse files
authored
[3.12] GH-117894: prevent aclose()/athrow() being re-used after StopIteration (GH-117851) (GH-118226)
(cherry picked from commit 7d369d4)
1 parent 33d1cff commit fc1732c

File tree

3 files changed

+70
-2
lines changed

3 files changed

+70
-2
lines changed

Lib/test/test_asyncgen.py

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -395,7 +395,11 @@ async def gen():
395395
yield 123
396396

397397
with self.assertWarns(DeprecationWarning):
398-
gen().athrow(GeneratorExit, GeneratorExit(), None)
398+
x = gen().athrow(GeneratorExit, GeneratorExit(), None)
399+
with self.assertRaises(GeneratorExit):
400+
x.send(None)
401+
del x
402+
gc_collect()
399403

400404
def test_async_gen_api_01(self):
401405
async def gen():
@@ -1653,6 +1657,62 @@ async def run():
16531657

16541658
self.loop.run_until_complete(run())
16551659

1660+
def test_async_gen_throw_same_aclose_coro_twice(self):
1661+
async def async_iterate():
1662+
yield 1
1663+
yield 2
1664+
1665+
it = async_iterate()
1666+
nxt = it.aclose()
1667+
with self.assertRaises(StopIteration):
1668+
nxt.throw(GeneratorExit)
1669+
1670+
with self.assertRaisesRegex(
1671+
RuntimeError,
1672+
r"cannot reuse already awaited aclose\(\)/athrow\(\)"
1673+
):
1674+
nxt.throw(GeneratorExit)
1675+
1676+
def test_async_gen_throw_custom_same_aclose_coro_twice(self):
1677+
async def async_iterate():
1678+
yield 1
1679+
yield 2
1680+
1681+
it = async_iterate()
1682+
1683+
class MyException(Exception):
1684+
pass
1685+
1686+
nxt = it.aclose()
1687+
with self.assertRaises(MyException):
1688+
nxt.throw(MyException)
1689+
1690+
with self.assertRaisesRegex(
1691+
RuntimeError,
1692+
r"cannot reuse already awaited aclose\(\)/athrow\(\)"
1693+
):
1694+
nxt.throw(MyException)
1695+
1696+
def test_async_gen_throw_custom_same_athrow_coro_twice(self):
1697+
async def async_iterate():
1698+
yield 1
1699+
yield 2
1700+
1701+
it = async_iterate()
1702+
1703+
class MyException(Exception):
1704+
pass
1705+
1706+
nxt = it.athrow(MyException)
1707+
with self.assertRaises(MyException):
1708+
nxt.throw(MyException)
1709+
1710+
with self.assertRaisesRegex(
1711+
RuntimeError,
1712+
r"cannot reuse already awaited aclose\(\)/athrow\(\)"
1713+
):
1714+
nxt.throw(MyException)
1715+
16561716
def test_async_gen_aclose_twice_with_different_coros(self):
16571717
# Regression test for https://bugs.python.org/issue39606
16581718
async def async_iterate():
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Prevent ``agen.aclose()`` objects being re-used after ``.throw()``.

Objects/genobject.c

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2223,7 +2223,11 @@ async_gen_athrow_throw(PyAsyncGenAThrow *o, PyObject *const *args, Py_ssize_t na
22232223

22242224
retval = gen_throw((PyGenObject*)o->agt_gen, args, nargs);
22252225
if (o->agt_args) {
2226-
return async_gen_unwrap_value(o->agt_gen, retval);
2226+
retval = async_gen_unwrap_value(o->agt_gen, retval);
2227+
if (retval == NULL) {
2228+
o->agt_state = AWAITABLE_STATE_CLOSED;
2229+
}
2230+
return retval;
22272231
} else {
22282232
/* aclose() mode */
22292233
if (retval && _PyAsyncGenWrappedValue_CheckExact(retval)) {
@@ -2233,6 +2237,9 @@ async_gen_athrow_throw(PyAsyncGenAThrow *o, PyObject *const *args, Py_ssize_t na
22332237
PyErr_SetString(PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG);
22342238
return NULL;
22352239
}
2240+
if (retval == NULL) {
2241+
o->agt_state = AWAITABLE_STATE_CLOSED;
2242+
}
22362243
if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration) ||
22372244
PyErr_ExceptionMatches(PyExc_GeneratorExit))
22382245
{

0 commit comments

Comments
 (0)