[Python-checkins] cpython (3.5): coroutines: Error when awaiting on coroutine that's being awaited

[yury.selivanov]

https://hg.python.org/cpython/rev/e16084e07761
changeset:   100400:e16084e07761
branch:      3.5
parent:      100398:bbc8cb86f05e
user:        Yury Selivanov <yselivanov at sprymix.com>
date:        Wed Mar 02 11:30:46 2016 -0500
summary:
  coroutines: Error when awaiting on coroutine that's being awaited

Issue #25888

files:
  Include/genobject.h         |   1 +
  Lib/test/test_coroutines.py |  18 ++++++++++++++++++
  Objects/genobject.c         |  12 ++++++------
  Python/ceval.c              |  15 +++++++++++++++
  4 files changed, 40 insertions(+), 6 deletions(-)


diff --git a/Include/genobject.h b/Include/genobject.h
--- a/Include/genobject.h
+++ b/Include/genobject.h
@@ -43,6 +43,7 @@
 PyAPI_FUNC(int) PyGen_NeedsFinalizing(PyGenObject *);
 PyAPI_FUNC(int) _PyGen_FetchStopIterationValue(PyObject **);
 PyObject *_PyGen_Send(PyGenObject *, PyObject *);
+PyObject *_PyGen_yf(PyGenObject *);
 PyAPI_FUNC(void) _PyGen_Finalize(PyObject *self);
 
 #ifndef Py_LIMITED_API
diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
--- a/Lib/test/test_coroutines.py
+++ b/Lib/test/test_coroutines.py
@@ -942,6 +942,24 @@
         with self.assertRaises(Marker):
             c.throw(ZeroDivisionError)
 
+    def test_await_15(self):
+        @types.coroutine
+        def nop():
+            yield
+
+        async def coroutine():
+            await nop()
+
+        async def waiter(coro):
+            await coro
+
+        coro = coroutine()
+        coro.send(None)
+
+        with self.assertRaisesRegex(RuntimeError,
+                                    "coroutine is being awaited already"):
+            waiter(coro).send(None)
+
     def test_with_1(self):
         class Manager:
             def __init__(self, name):
diff --git a/Objects/genobject.c b/Objects/genobject.c
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -267,8 +267,8 @@
     return 0;
 }
 
-static PyObject *
-gen_yf(PyGenObject *gen)
+PyObject *
+_PyGen_yf(PyGenObject *gen)
 {
     PyObject *yf = NULL;
     PyFrameObject *f = gen->gi_frame;
@@ -290,7 +290,7 @@
 gen_close(PyGenObject *gen, PyObject *args)
 {
     PyObject *retval;
-    PyObject *yf = gen_yf(gen);
+    PyObject *yf = _PyGen_yf(gen);
     int err = 0;
 
     if (yf) {
@@ -330,7 +330,7 @@
     PyObject *typ;
     PyObject *tb = NULL;
     PyObject *val = NULL;
-    PyObject *yf = gen_yf(gen);
+    PyObject *yf = _PyGen_yf(gen);
     _Py_IDENTIFIER(throw);
 
     if (!PyArg_UnpackTuple(args, "throw", 1, 3, &typ, &val, &tb))
@@ -564,7 +564,7 @@
 static PyObject *
 gen_getyieldfrom(PyGenObject *gen)
 {
-    PyObject *yf = gen_yf(gen);
+    PyObject *yf = _PyGen_yf(gen);
     if (yf == NULL)
         Py_RETURN_NONE;
     return yf;
@@ -799,7 +799,7 @@
 static PyObject *
 coro_get_cr_await(PyCoroObject *coro)
 {
-    PyObject *yf = gen_yf((PyGenObject *) coro);
+    PyObject *yf = _PyGen_yf((PyGenObject *) coro);
     if (yf == NULL)
         Py_RETURN_NONE;
     return yf;
diff --git a/Python/ceval.c b/Python/ceval.c
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -2021,6 +2021,21 @@
 
             Py_DECREF(iterable);
 
+            if (iter != NULL && PyCoro_CheckExact(iter)) {
+                PyObject *yf = _PyGen_yf((PyGenObject*)iter);
+                if (yf != NULL) {
+                    /* `iter` is a coroutine object that is being
+                       awaited, `yf` is a pointer to the current awaitable
+                       being awaited on. */
+                    Py_DECREF(yf);
+                    Py_CLEAR(iter);
+                    PyErr_SetString(
+                        PyExc_RuntimeError,
+                        "coroutine is being awaited already");
+                    /* The code below jumps to `error` if `iter` is NULL. */
+                }
+            }
+
             SET_TOP(iter); /* Even if it's NULL */
 
             if (iter == NULL) {

-- 
Repository URL: https://hg.python.org/cpython