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:
  https://github.com/python-trio/trio/pull/1396


https://bugs.python.org/issue39606
This commit is contained in:
Nathaniel J. Smith 2020-02-13 00:15:38 -08:00 committed by GitHub
parent 7514f4f625
commit 925dc7fb1d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 41 additions and 6 deletions

View file

@ -1797,16 +1797,22 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
PyFrameObject *f = gen->gi_frame;
PyObject *retval;
if (f == NULL || f->f_stacktop == NULL ||
o->agt_state == AWAITABLE_STATE_CLOSED) {
if (o->agt_state == AWAITABLE_STATE_CLOSED) {
PyErr_SetString(
PyExc_RuntimeError,
"cannot reuse already awaited aclose()/athrow()");
return NULL;
}
if (f == NULL || f->f_stacktop == NULL) {
o->agt_state = AWAITABLE_STATE_CLOSED;
PyErr_SetNone(PyExc_StopIteration);
return NULL;
}
if (o->agt_state == AWAITABLE_STATE_INIT) {
if (o->agt_gen->ag_running_async) {
o->agt_state = AWAITABLE_STATE_CLOSED;
if (o->agt_args == NULL) {
PyErr_SetString(
PyExc_RuntimeError,
@ -1878,7 +1884,6 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
/* aclose() mode */
if (retval) {
if (_PyAsyncGenWrappedValue_CheckExact(retval)) {
o->agt_gen->ag_running_async = 0;
Py_DECREF(retval);
goto yield_close;
}
@ -1893,16 +1898,17 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
yield_close:
o->agt_gen->ag_running_async = 0;
o->agt_state = AWAITABLE_STATE_CLOSED;
PyErr_SetString(
PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG);
return NULL;
check_error:
o->agt_gen->ag_running_async = 0;
o->agt_state = AWAITABLE_STATE_CLOSED;
if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration) ||
PyErr_ExceptionMatches(PyExc_GeneratorExit))
{
o->agt_state = AWAITABLE_STATE_CLOSED;
if (o->agt_args == NULL) {
/* when aclose() is called we don't want to propagate
StopAsyncIteration or GeneratorExit; just raise
@ -1936,6 +1942,7 @@ async_gen_athrow_throw(PyAsyncGenAThrow *o, PyObject *args)
/* aclose() mode */
if (retval && _PyAsyncGenWrappedValue_CheckExact(retval)) {
o->agt_gen->ag_running_async = 0;
o->agt_state = AWAITABLE_STATE_CLOSED;
Py_DECREF(retval);
PyErr_SetString(PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG);
return NULL;