diff --git a/src/iter.cpp b/src/iter.cpp index 4cc1e796..eabc4790 100644 --- a/src/iter.cpp +++ b/src/iter.cpp @@ -49,7 +49,15 @@ namespace pkpy{ for(PyObject* obj: s_backup) frame._s->push(obj); s_backup.clear(); vm->callstack.push(std::move(frame)); - PyObject* ret = vm->_run_top_frame(); + + PyObject* ret; + try{ + ret = vm->_run_top_frame(); + }catch(...){ + state = 2; // end this generator immediately when an exception is thrown + throw; + } + if(ret == PY_OP_YIELD){ // backup the context frame = std::move(vm->callstack.top()); diff --git a/tests/45_yield.py b/tests/45_yield.py index 40e3cdcc..2d113fba 100644 --- a/tests/45_yield.py +++ b/tests/45_yield.py @@ -45,4 +45,23 @@ def f(n): yield j t = f(3) -assert list(t) == [0, 1, 0, 2, 0, 1] \ No newline at end of file +assert list(t) == [0, 1, 0, 2, 0, 1] + +def f(n): + for i in range(n): + if i == n-1: + raise ValueError + yield i + +t = f(3) +t = iter(t) +assert next(t) == 0 +assert next(t) == 1 + +try: + next(t) + exit(1) +except ValueError: + pass + +assert next(t) == StopIteration \ No newline at end of file