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

4 files changed, 280 insertions, 22 deletions

diff --git a/Objects/abstract.c b/Objects/abstract.c
index a44bec6..4d73a3b 100644
--- a/Objects/abstract.c
+++ b/Objects/abstract.c
@@ -2267,7 +2267,6 @@ PyObject_CallMethod(PyObject *o, char *name, char *format, ...)
 
     func = PyObject_GetAttrString(o, name);
     if (func == NULL) {
-        PyErr_SetString(PyExc_AttributeError, name);
         return 0;
     }
 
@@ -2311,7 +2310,6 @@ _PyObject_CallMethod_SizeT(PyObject *o, char *name, char *format, ...)
 
     func = PyObject_GetAttrString(o, name);
     if (func == NULL) {
-        PyErr_SetString(PyExc_AttributeError, name);
         return 0;
     }
     va_start(va, format);
diff --git a/Objects/exceptions.c b/Objects/exceptions.c
index 3318115..a529a3c 100644
--- a/Objects/exceptions.c
+++ b/Objects/exceptions.c
@@ -487,8 +487,70 @@ SimpleExtendsException(PyExc_Exception, TypeError,
 /*
  *    StopIteration extends Exception
  */
-SimpleExtendsException(PyExc_Exception, StopIteration,
-                       "Signal the end from iterator.__next__().");
+
+static PyMemberDef StopIteration_members[] = {
+    {"value", T_OBJECT, offsetof(PyStopIterationObject, value), 0,
+        PyDoc_STR("generator return value")},
+    {NULL}  /* Sentinel */
+};
+
+static int
+StopIteration_init(PyStopIterationObject *self, PyObject *args, PyObject *kwds)
+{
+    Py_ssize_t size = PyTuple_GET_SIZE(args);
+    PyObject *value;
+
+    if (BaseException_init((PyBaseExceptionObject *)self, args, kwds) == -1)
+        return -1;
+    Py_CLEAR(self->value);
+    if (size > 0)
+        value = PyTuple_GET_ITEM(args, 0);
+    else
+        value = Py_None;
+    Py_INCREF(value);
+    self->value = value;
+    return 0;
+}
+
+static int
+StopIteration_clear(PyStopIterationObject *self)
+{
+    Py_CLEAR(self->value);
+    return BaseException_clear((PyBaseExceptionObject *)self);
+}
+
+static void
+StopIteration_dealloc(PyStopIterationObject *self)
+{
+    _PyObject_GC_UNTRACK(self);
+    StopIteration_clear(self);
+    Py_TYPE(self)->tp_free((PyObject *)self);
+}
+
+static int
+StopIteration_traverse(PyStopIterationObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(self->value);
+    return BaseException_traverse((PyBaseExceptionObject *)self, visit, arg);
+}
+
+PyObject *
+PyStopIteration_Create(PyObject *value)
+{
+    return PyObject_CallFunctionObjArgs(PyExc_StopIteration, value, NULL);
+}
+
+ComplexExtendsException(
+    PyExc_Exception,       /* base */
+    StopIteration,         /* name */
+    StopIteration,         /* prefix for *_init, etc */
+    0,                     /* new */
+    0,                     /* methods */
+    StopIteration_members, /* members */
+    0,                     /* getset */
+    0,                     /* str */
+    "Signal the end from iterator.__next__()."
+);
 
 
 /*
diff --git a/Objects/frameobject.c b/Objects/frameobject.c
index 10fb8b3..9b05b9d 100644
--- a/Objects/frameobject.c
+++ b/Objects/frameobject.c
@@ -15,11 +15,12 @@
 #define OFF(x) offsetof(PyFrameObject, x)
 
 static PyMemberDef frame_memberlist[] = {
-    {"f_back",          T_OBJECT,       OFF(f_back),    READONLY},
-    {"f_code",          T_OBJECT,       OFF(f_code),    READONLY},
-    {"f_builtins",      T_OBJECT,       OFF(f_builtins),READONLY},
-    {"f_globals",       T_OBJECT,       OFF(f_globals), READONLY},
-    {"f_lasti",         T_INT,          OFF(f_lasti),   READONLY},
+    {"f_back",          T_OBJECT,       OFF(f_back),      READONLY},
+    {"f_code",          T_OBJECT,       OFF(f_code),      READONLY},
+    {"f_builtins",      T_OBJECT,       OFF(f_builtins),  READONLY},
+    {"f_globals",       T_OBJECT,       OFF(f_globals),   READONLY},
+    {"f_lasti",         T_INT,          OFF(f_lasti),     READONLY},
+    {"f_yieldfrom",     T_OBJECT,       OFF(f_yieldfrom), READONLY},
     {NULL}      /* Sentinel */
 };
 
@@ -444,6 +445,7 @@ frame_dealloc(PyFrameObject *f)
     Py_CLEAR(f->f_exc_type);
     Py_CLEAR(f->f_exc_value);
     Py_CLEAR(f->f_exc_traceback);
+    Py_CLEAR(f->f_yieldfrom);
 
     co = f->f_code;
     if (co->co_zombieframe == NULL)
@@ -475,6 +477,7 @@ frame_traverse(PyFrameObject *f, visitproc visit, void *arg)
     Py_VISIT(f->f_exc_type);
     Py_VISIT(f->f_exc_value);
     Py_VISIT(f->f_exc_traceback);
+    Py_VISIT(f->f_yieldfrom);
 
     /* locals */
     slots = f->f_code->co_nlocals + PyTuple_GET_SIZE(f->f_code->co_cellvars) + PyTuple_GET_SIZE(f->f_code->co_freevars);
@@ -508,6 +511,7 @@ frame_clear(PyFrameObject *f)
     Py_CLEAR(f->f_exc_value);
     Py_CLEAR(f->f_exc_traceback);
     Py_CLEAR(f->f_trace);
+    Py_CLEAR(f->f_yieldfrom);
 
     /* locals */
     slots = f->f_code->co_nlocals + PyTuple_GET_SIZE(f->f_code->co_cellvars) + PyTuple_GET_SIZE(f->f_code->co_freevars);
@@ -711,6 +715,7 @@ PyFrame_New(PyThreadState *tstate, PyCodeObject *code, PyObject *globals,
     f->f_lasti = -1;
     f->f_lineno = code->co_firstlineno;
     f->f_iblock = 0;
+    f->f_yieldfrom = NULL;
 
     _PyObject_GC_TRACK(f);
     return f;
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)