From a96e06db77dcbd3433d39761ddb4615d7d96284a Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Tue, 21 Jan 2020 00:49:30 +0200 Subject: [PATCH] bpo-39386: Prevent double awaiting of async iterator (GH-18081) --- Lib/test/test_asyncgen.py | 36 +++++++++++++++++++ .../2020-01-20-21-40-57.bpo-39386.ULqD8t.rst | 1 + Objects/genobject.c | 16 ++++++--- 3 files changed, 49 insertions(+), 4 deletions(-) create mode 100644 Misc/NEWS.d/next/Core and Builtins/2020-01-20-21-40-57.bpo-39386.ULqD8t.rst diff --git a/Lib/test/test_asyncgen.py b/Lib/test/test_asyncgen.py index 58d8aee19ad..24b20bec2b2 100644 --- a/Lib/test/test_asyncgen.py +++ b/Lib/test/test_asyncgen.py @@ -1128,6 +1128,42 @@ class AsyncGenAsyncioTest(unittest.TestCase): self.assertEqual([], messages) + def test_async_gen_await_anext_twice(self): + async def async_iterate(): + yield 1 + yield 2 + + async def run(): + it = async_iterate() + nxt = it.__anext__() + await nxt + with self.assertRaisesRegex( + RuntimeError, + r"cannot reuse already awaited __anext__\(\)/asend\(\)" + ): + await nxt + + await it.aclose() # prevent unfinished iterator warning + + self.loop.run_until_complete(run()) + + def test_async_gen_await_aclose_twice(self): + async def async_iterate(): + yield 1 + yield 2 + + async def run(): + it = async_iterate() + nxt = it.aclose() + await nxt + with self.assertRaisesRegex( + RuntimeError, + r"cannot reuse already awaited aclose\(\)/athrow\(\)" + ): + await nxt + + self.loop.run_until_complete(run()) + if __name__ == "__main__": unittest.main() diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-01-20-21-40-57.bpo-39386.ULqD8t.rst b/Misc/NEWS.d/next/Core and Builtins/2020-01-20-21-40-57.bpo-39386.ULqD8t.rst new file mode 100644 index 00000000000..f24e1f4e8a1 --- /dev/null +++ b/Misc/NEWS.d/next/Core and Builtins/2020-01-20-21-40-57.bpo-39386.ULqD8t.rst @@ -0,0 +1 @@ +Prevent double awaiting of async iterator. diff --git a/Objects/genobject.c b/Objects/genobject.c index c5fe9999af2..652c2903dd2 100644 --- a/Objects/genobject.c +++ b/Objects/genobject.c @@ -1518,7 +1518,9 @@ async_gen_asend_send(PyAsyncGenASend *o, PyObject *arg) PyObject *result; if (o->ags_state == AWAITABLE_STATE_CLOSED) { - PyErr_SetNone(PyExc_StopIteration); + PyErr_SetString( + PyExc_RuntimeError, + "cannot reuse already awaited __anext__()/asend()"); return NULL; } @@ -1561,7 +1563,9 @@ async_gen_asend_throw(PyAsyncGenASend *o, PyObject *args) PyObject *result; if (o->ags_state == AWAITABLE_STATE_CLOSED) { - PyErr_SetNone(PyExc_StopIteration); + PyErr_SetString( + PyExc_RuntimeError, + "cannot reuse already awaited __anext__()/asend()"); return NULL; } @@ -1795,7 +1799,9 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg) if (f == NULL || f->f_stacktop == NULL || o->agt_state == AWAITABLE_STATE_CLOSED) { - PyErr_SetNone(PyExc_StopIteration); + PyErr_SetString( + PyExc_RuntimeError, + "cannot reuse already awaited aclose()/athrow()"); return NULL; } @@ -1917,7 +1923,9 @@ async_gen_athrow_throw(PyAsyncGenAThrow *o, PyObject *args) PyObject *retval; if (o->agt_state == AWAITABLE_STATE_CLOSED) { - PyErr_SetNone(PyExc_StopIteration); + PyErr_SetString( + PyExc_RuntimeError, + "cannot reuse already awaited aclose()/athrow()"); return NULL; }