From 099a78fe6d0a4c20126497940529d2b1fd41c5de Mon Sep 17 00:00:00 2001 From: Benjamin Peterson Date: Wed, 7 Mar 2012 17:57:04 -0600 Subject: make delegating generators say they running (closes #14220) --- Lib/test/test_pep380.py | 72 +++++++++++++++++++++++++++++++++++++++++++++++++ Misc/NEWS | 3 +++ Objects/genobject.c | 51 +++++++++++++++++++++++++---------- 3 files changed, 112 insertions(+), 14 deletions(-) diff --git a/Lib/test/test_pep380.py b/Lib/test/test_pep380.py index 3a1db29..bdcfacd 100644 --- a/Lib/test/test_pep380.py +++ b/Lib/test/test_pep380.py @@ -847,6 +847,78 @@ class TestPEP380Operation(unittest.TestCase): yield from () self.assertRaises(StopIteration, next, g()) + def test_delegating_generators_claim_to_be_running(self): + # Check with basic iteration + def one(): + yield 0 + yield from two() + yield 3 + def two(): + yield 1 + try: + yield from g1 + except ValueError: + pass + yield 2 + g1 = one() + self.assertEqual(list(g1), [0, 1, 2, 3]) + # Check with send + g1 = one() + res = [next(g1)] + try: + while True: + res.append(g1.send(42)) + except StopIteration: + pass + self.assertEqual(res, [0, 1, 2, 3]) + # Check with throw + class MyErr(Exception): + pass + def one(): + try: + yield 0 + except MyErr: + pass + yield from two() + try: + yield 3 + except MyErr: + pass + def two(): + try: + yield 1 + except MyErr: + pass + try: + yield from g1 + except ValueError: + pass + try: + yield 2 + except MyErr: + pass + g1 = one() + res = [next(g1)] + try: + while True: + res.append(g1.throw(MyErr)) + except StopIteration: + pass + # Check with close + class MyIt(object): + def __iter__(self): + return self + def __next__(self): + return 42 + def close(self_): + self.assertTrue(g1.gi_running) + self.assertRaises(ValueError, next, g1) + def one(): + yield from MyIt() + g1 = one() + next(g1) + g1.close() + def test_main(): from test import support diff --git a/Misc/NEWS b/Misc/NEWS index 7d64411..a8769af 100644 --- a/Misc/NEWS +++ b/Misc/NEWS @@ -13,6 +13,9 @@ Core and Builtins - Issue #14205: dict lookup raises a RuntimeError if the dict is modified during a lookup. +- Issue #14220: When a generator is delegating to another iterator with the + yield from syntax, it needs to have its ``gi_running`` flag set to True. + Library ------- diff --git a/Objects/genobject.c b/Objects/genobject.c index b32d9b6..e41779e 100644 --- a/Objects/genobject.c +++ b/Objects/genobject.c @@ -41,6 +41,15 @@ gen_dealloc(PyGenObject *gen) PyObject_GC_Del(gen); } +static int +gen_running(PyGenObject *gen) +{ + if (gen->gi_running) { + PyErr_SetString(PyExc_ValueError, "generator already executing"); + return 1; + } + return 0; +} static PyObject * gen_send_ex(PyGenObject *gen, PyObject *arg, int exc) @@ -49,11 +58,7 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc) PyFrameObject *f = gen->gi_frame; PyObject *result; - if (gen->gi_running) { - PyErr_SetString(PyExc_ValueError, - "generator already executing"); - return NULL; - } + assert(!gen->gi_running); if (f==NULL || f->f_stacktop == NULL) { /* Only set exception if called from send() */ if (arg && !exc) @@ -137,12 +142,15 @@ gen_send(PyGenObject *gen, PyObject *arg) int exc = 0; PyObject *ret; PyObject *yf = gen->gi_frame ? gen->gi_frame->f_yieldfrom : NULL; + if (gen_running(gen)) + return 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); + gen->gi_running = 1; if (PyGen_CheckExact(yf)) { ret = gen_send((PyGenObject *)yf, arg); } else { @@ -151,6 +159,7 @@ gen_send(PyGenObject *gen, PyObject *arg) else ret = PyObject_CallMethod(yf, "send", "O", arg); } + gen->gi_running = 0; if (ret) { Py_DECREF(yf); goto done; @@ -177,17 +186,19 @@ PyDoc_STRVAR(close_doc, */ static int -gen_close_iter(PyObject *yf) +gen_close_iter(PyGenObject *gen, PyObject *yf) { PyObject *retval = NULL; + int err = 0; if (PyGen_CheckExact(yf)) { retval = gen_close((PyGenObject *)yf, NULL); - if (retval == NULL) { - return -1; - } + if (!retval) + err = -1; } else { - PyObject *meth = PyObject_GetAttrString(yf, "close"); + PyObject *meth; + gen->gi_running = 1; + meth = PyObject_GetAttrString(yf, "close"); if (meth == NULL) { if (!PyErr_ExceptionMatches(PyExc_AttributeError)) { PyErr_WriteUnraisable(yf); @@ -197,11 +208,12 @@ gen_close_iter(PyObject *yf) retval = PyObject_CallFunction(meth, ""); Py_DECREF(meth); if (!retval) - return -1; + err = -1; } + gen->gi_running = 0; } Py_XDECREF(retval); - return 0; + return err; } static PyObject * @@ -211,9 +223,11 @@ gen_close(PyGenObject *gen, PyObject *args) PyObject *yf = gen->gi_frame ? gen->gi_frame->f_yieldfrom : NULL; int err = 0; + if (gen_running(gen)) + return NULL; if (yf) { Py_INCREF(yf); - err = gen_close_iter(yf); + err = gen_close_iter(gen, yf); gen_undelegate(gen); Py_DECREF(yf); } @@ -314,18 +328,22 @@ gen_throw(PyGenObject *gen, PyObject *args) if (!PyArg_UnpackTuple(args, "throw", 1, 3, &typ, &val, &tb)) return NULL; + if (gen_running(gen)) + return NULL; + if (yf) { PyObject *ret; int err; Py_INCREF(yf); if (PyErr_GivenExceptionMatches(typ, PyExc_GeneratorExit)) { - err = gen_close_iter(yf); + err = gen_close_iter(gen, yf); Py_DECREF(yf); gen_undelegate(gen); if (err < 0) return gen_send_ex(gen, Py_None, 1); goto throw_here; } + gen->gi_running = 1; if (PyGen_CheckExact(yf)) { ret = gen_throw((PyGenObject *)yf, args); } else { @@ -343,6 +361,7 @@ gen_throw(PyGenObject *gen, PyObject *args) ret = PyObject_CallObject(meth, args); Py_DECREF(meth); } + gen->gi_running = 0; Py_DECREF(yf); if (!ret) { PyObject *val; @@ -423,10 +442,14 @@ gen_iternext(PyGenObject *gen) PyObject *ret; int exc = 0; PyObject *yf = gen->gi_frame ? gen->gi_frame->f_yieldfrom : NULL; + if (gen_running(gen)) + return NULL; if (yf) { Py_INCREF(yf); /* ceval.c ensures that yf is an iterator */ + gen->gi_running = 1; ret = Py_TYPE(yf)->tp_iternext(yf); + gen->gi_running = 0; if (ret) { Py_DECREF(yf); return ret; -- cgit v0.12