Implement PEP 380 - 'yield from' (closes #11682)
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 @@
/* 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 @@
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 @@
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 @@
_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 @@
* 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 @@
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 @@
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 @@
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)