Implement PEP 380 - 'yield from' (closes #11682)

1 files changed, 206 insertions, 13 deletions

diff --git a/Objects/genobject.c b/Objects/genobject.c
index c6612e2..20c926b 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -5,6 +5,9 @@
 #include "structmember.h"
 #include "opcode.h"
 
+static PyObject *gen_close(PyGenObject *gen, PyObject *args);
+static void gen_undelegate(PyGenObject *gen);
+
 static int
 gen_traverse(PyGenObject *gen, visitproc visit, void *arg)
 {
@@ -90,12 +93,18 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc)
 
     /* If the generator just returned (as opposed to yielding), signal
      * that the generator is exhausted. */
-    if (result == Py_None && f->f_stacktop == NULL) {
-        Py_DECREF(result);
-        result = NULL;
-        /* Set exception if not called by gen_iternext() */
-        if (arg)
+    if (result && f->f_stacktop == NULL) {
+        if (result == Py_None) {
+            /* Delay exception instantiation if we can */
             PyErr_SetNone(PyExc_StopIteration);
+        } else {
+            PyObject *e = PyStopIteration_Create(result);
+            if (e != NULL) {
+                PyErr_SetObject(PyExc_StopIteration, e);
+                Py_DECREF(e);
+            }
+        }
+        Py_CLEAR(result);
     }
 
     if (!result || f->f_stacktop == NULL) {
@@ -111,8 +120,8 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc)
         Py_XDECREF(t);
         Py_XDECREF(v);
         Py_XDECREF(tb);
-        Py_DECREF(f);
         gen->gi_frame = NULL;
+        Py_DECREF(f);
     }
 
     return result;
@@ -125,17 +134,91 @@ return next yielded value or raise StopIteration.");
 static PyObject *
 gen_send(PyGenObject *gen, PyObject *arg)
 {
-    return gen_send_ex(gen, arg, 0);
+    int exc = 0;
+    PyObject *ret;
+    PyObject *yf = gen->gi_frame ? gen->gi_frame->f_yieldfrom : NULL;
+    /* XXX (ncoghlan): Are the incref/decref on arg and yf strictly needed?
+     *                 Or would it be valid to rely on borrowed references?
+     */
+    Py_INCREF(arg);
+    if (yf) {
+        Py_INCREF(yf);
+        if (PyGen_CheckExact(yf)) {
+            ret = gen_send((PyGenObject *)yf, arg);
+        } else {
+            if (arg == Py_None)
+                ret = PyIter_Next(yf);
+            else
+                ret = PyObject_CallMethod(yf, "send", "O", arg);
+        }
+        if (ret) {
+            Py_DECREF(yf);
+            goto done;
+        }
+        gen_undelegate(gen);
+        Py_CLEAR(arg);
+        if (PyGen_FetchStopIterationValue(&arg) < 0) {
+            exc = 1;
+        }
+        Py_DECREF(yf);
+    }
+    ret = gen_send_ex(gen, arg, exc);
+done:
+    Py_XDECREF(arg);
+    return ret;
 }
 
 PyDoc_STRVAR(close_doc,
 "close(arg) -> raise GeneratorExit inside generator.");
 
+/*
+ *   This helper function is used by gen_close and gen_throw to
+ *   close a subiterator being delegated to by yield-from.
+ */
+
+static int
+gen_close_iter(PyObject *yf)
+{
+    PyObject *retval = NULL;
+    
+    if (PyGen_CheckExact(yf)) {
+        retval = gen_close((PyGenObject *)yf, NULL);
+        if (retval == NULL) {
+            return -1;
+        }
+    } else {
+        PyObject *meth = PyObject_GetAttrString(yf, "close");
+        if (meth == NULL) {
+            if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
+                PyErr_WriteUnraisable(yf);
+            }
+            PyErr_Clear();
+        } else {
+            retval = PyObject_CallFunction(meth, "");
+            Py_DECREF(meth);
+            if (!retval)
+                return -1;
+        }
+    }
+    Py_XDECREF(retval);
+    return 0;
+}       
+
 static PyObject *
 gen_close(PyGenObject *gen, PyObject *args)
 {
     PyObject *retval;
-    PyErr_SetNone(PyExc_GeneratorExit);
+    PyObject *yf = gen->gi_frame ? gen->gi_frame->f_yieldfrom : NULL;
+    int err = 0;
+
+    if (yf) {
+        Py_INCREF(yf);
+        err = gen_close_iter(yf);
+        gen_undelegate(gen);
+        Py_DECREF(yf);
+    }
+    if (err == 0)
+        PyErr_SetNone(PyExc_GeneratorExit);
     retval = gen_send_ex(gen, Py_None, 1);
     if (retval) {
         Py_DECREF(retval);
@@ -196,7 +279,7 @@ gen_del(PyObject *self)
         _Py_NewReference(self);
         self->ob_refcnt = refcnt;
     }
-    assert(PyType_IS_GC(self->ob_type) &&
+    assert(PyType_IS_GC(Py_TYPE(self)) &&
            _Py_AS_GC(self)->gc.gc_refs != _PyGC_REFS_UNTRACKED);
 
     /* If Py_REF_DEBUG, _Py_NewReference bumped _Py_RefTotal, so
@@ -209,8 +292,8 @@ gen_del(PyObject *self)
      * undone.
      */
 #ifdef COUNT_ALLOCS
-    --self->ob_type->tp_frees;
-    --self->ob_type->tp_allocs;
+    --(Py_TYPE(self)->tp_frees);
+    --(Py_TYPE(self)->tp_allocs);
 #endif
 }
 
@@ -226,10 +309,55 @@ gen_throw(PyGenObject *gen, PyObject *args)
     PyObject *typ;
     PyObject *tb = NULL;
     PyObject *val = NULL;
+    PyObject *yf = gen->gi_frame ? gen->gi_frame->f_yieldfrom : NULL;
 
     if (!PyArg_UnpackTuple(args, "throw", 1, 3, &typ, &val, &tb))
         return NULL;
 
+    if (yf) {
+        PyObject *ret;
+        int err;
+        Py_INCREF(yf);
+        if (PyErr_GivenExceptionMatches(typ, PyExc_GeneratorExit)) {
+            err = gen_close_iter(yf);
+            Py_DECREF(yf);
+            gen_undelegate(gen);
+            if (err < 0)
+                return gen_send_ex(gen, Py_None, 1);
+            goto throw_here;
+        }
+        if (PyGen_CheckExact(yf)) {
+            ret = gen_throw((PyGenObject *)yf, args);
+        } else {
+            PyObject *meth = PyObject_GetAttrString(yf, "throw");
+            if (meth == NULL) {
+                if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
+                    Py_DECREF(yf);
+                    return NULL;
+                }
+                PyErr_Clear();
+                Py_DECREF(yf);
+                gen_undelegate(gen);
+                goto throw_here;
+            }
+            ret = PyObject_CallObject(meth, args);
+            Py_DECREF(meth);
+        }
+        Py_DECREF(yf);
+        if (!ret) {
+            PyObject *val;
+            gen_undelegate(gen);
+            if (PyGen_FetchStopIterationValue(&val) == 0) {
+                ret = gen_send_ex(gen, val, 0);
+                Py_DECREF(val);
+            } else {
+                ret = gen_send_ex(gen, Py_None, 1);
+            }
+        }
+        return ret;
+    }
+
+throw_here:
     /* First, check the traceback argument, replacing None with
        NULL. */
     if (tb == Py_None) {
@@ -272,7 +400,7 @@ gen_throw(PyGenObject *gen, PyObject *args)
         PyErr_Format(PyExc_TypeError,
                      "exceptions must be classes or instances "
                      "deriving from BaseException, not %s",
-                     typ->ob_type->tp_name);
+                     Py_TYPE(typ)->tp_name);
             goto failed_throw;
     }
 
@@ -291,9 +419,74 @@ failed_throw:
 static PyObject *
 gen_iternext(PyGenObject *gen)
 {
-    return gen_send_ex(gen, NULL, 0);
+    PyObject *val = NULL;
+    PyObject *ret;
+    int exc = 0;
+    PyObject *yf = gen->gi_frame ? gen->gi_frame->f_yieldfrom : NULL;
+    if (yf) {
+        Py_INCREF(yf);
+        /* ceval.c ensures that yf is an iterator */
+        ret = Py_TYPE(yf)->tp_iternext(yf);
+        if (ret) {
+            Py_DECREF(yf);
+            return ret;
+        }
+        gen_undelegate(gen);
+        if (PyGen_FetchStopIterationValue(&val) < 0)
+            exc = 1;
+        Py_DECREF(yf);
+    }
+    ret = gen_send_ex(gen, val, exc);
+    Py_XDECREF(val);
+    return ret;
+}
+
+/*
+ *   In certain recursive situations, a generator may lose its frame
+ *   before we get a chance to clear f_yieldfrom, so we use this
+ *   helper function.
+ */
+
+static void
+gen_undelegate(PyGenObject *gen) {
+    if (gen->gi_frame) {
+        Py_XDECREF(gen->gi_frame->f_yieldfrom);
+        gen->gi_frame->f_yieldfrom = NULL;
+    }
 }
 
+/*
+ *   If StopIteration exception is set, fetches its 'value'
+ *   attribute if any, otherwise sets pvalue to None.
+ *
+ *   Returns 0 if no exception or StopIteration is set.
+ *   If any other exception is set, returns -1 and leaves
+ *   pvalue unchanged.
+ */
+
+int
+PyGen_FetchStopIterationValue(PyObject **pvalue) {
+    PyObject *et, *ev, *tb;
+    PyObject *value = NULL;
+    
+    if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
+        PyErr_Fetch(&et, &ev, &tb);
+        Py_XDECREF(et);
+        Py_XDECREF(tb);
+        if (ev) {
+            value = ((PyStopIterationObject *)ev)->value;
+            Py_DECREF(ev);
+        }
+    } else if (PyErr_Occurred()) {
+        return -1;
+    }
+    if (value == NULL) {
+        value = Py_None;
+    }
+    Py_INCREF(value);
+    *pvalue = value;
+    return 0;
+}
 
 static PyObject *
 gen_repr(PyGenObject *gen)