From ca7fb1578778a4ce10ac89a563f22bc87c8b1e8b Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Thu, 7 Sep 2023 00:35:47 +0800
Subject: [PATCH] kill generator on body raise

---
 src/iter.cpp      | 10 +++++++++-
 tests/45_yield.py | 21 ++++++++++++++++++++-
 2 files changed, 29 insertions(+), 2 deletions(-)

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