Move exception state cleanup into generator body code to allow a distinction...

Move exception state cleanup into generator body code to allow a distinction between normal yields and yields in except blocks. They require different handling according to what CPython does.

Cython/Compiler/Code.pxd

@@ -33,6 +33,7 @@ cdef class FunctionState:
     cdef public object return_from_error_cleanup_label # not used in __init__ ?
 
     cdef public object exc_vars
+    cdef public object current_except
     cdef public bint in_try_finally
     cdef public bint can_trace
     cdef public bint gil_owned

Cython/Compiler/Code.py

@@ -602,6 +602,7 @@ class FunctionState(object):
 
         self.in_try_finally = 0
         self.exc_vars = None
+        self.current_except = None
         self.can_trace = False
         self.gil_owned = True
 

Cython/Compiler/ExprNodes.py

@@ -9469,6 +9469,13 @@ class YieldExprNode(ExprNode):
                                   nogil=not code.funcstate.gil_owned)
         code.put_finish_refcount_context()
 
+        if code.funcstate.current_except is not None:
+            # inside of an except block => save away currently handled exception
+            code.putln("__Pyx_Coroutine_SwapException(%s);" % Naming.generator_cname)
+        else:
+            # no exceptions being handled => restore exception state of caller
+            code.putln("__Pyx_Coroutine_ResetAndClearException(%s);" % Naming.generator_cname)
+
         code.putln("/* return from %sgenerator, %sing value */" % (
             'async ' if self.in_async_gen else '',
             'await' if self.is_await else 'yield'))

Cython/Compiler/Nodes.py

@@ -4081,10 +4081,12 @@ class GeneratorBodyDefNode(DefNode):
         self.generate_function_header(code)
         closure_init_code = code.insertion_point()
         # ----- Local variables
+        code.putln("__Pyx_PyThreadState_declare")
         code.putln("PyObject *%s = NULL;" % Naming.retval_cname)
         tempvardecl_code = code.insertion_point()
         code.put_declare_refcount_context()
         code.put_setup_refcount_context(self.entry.name or self.entry.qualified_name)
+        code.putln("__Pyx_PyThreadState_assign")
         profile = code.globalstate.directives['profile']
         linetrace = code.globalstate.directives['linetrace']
         if profile or linetrace:
@@ -4157,6 +4159,7 @@ class GeneratorBodyDefNode(DefNode):
             code.put_xgiveref(Naming.retval_cname)
         else:
             code.put_xdecref_clear(Naming.retval_cname, py_object_type)
+        code.putln("__Pyx_Coroutine_ResetAndClearException(%s);" % Naming.generator_cname)
         code.putln('%s->resume_label = -1;' % Naming.generator_cname)
         # clean up as early as possible to help breaking any reference cycles
         code.putln('__Pyx_Coroutine_clear((PyObject*)%s);' % Naming.generator_cname)
@@ -6696,6 +6699,7 @@ class TryExceptStatNode(StatNode):
     #  else_clause      StatNode or None
 
     child_attrs = ["body", "except_clauses", "else_clause"]
+    in_generator = False
 
     def analyse_declarations(self, env):
         self.body.analyse_declarations(env)
@@ -6755,6 +6759,7 @@ class TryExceptStatNode(StatNode):
         if can_raise:
             # inject code before the try block to save away the exception state
             code.globalstate.use_utility_code(reset_exception_utility_code)
+            if not self.in_generator:
                 save_exc.putln("__Pyx_PyThreadState_declare")
                 save_exc.putln("__Pyx_PyThreadState_assign")
             save_exc.putln("__Pyx_ExceptionSave(%s);" % (
@@ -6794,11 +6799,16 @@ class TryExceptStatNode(StatNode):
                     code.put_xdecref_clear(var, py_object_type)
                 code.put_goto(try_end_label)
             code.put_label(our_error_label)
-            code.putln("__Pyx_PyThreadState_assign")  # re-assign in case a generator yielded
             for temp_name, temp_type in temps_to_clean_up:
                 code.put_xdecref_clear(temp_name, temp_type)
+
+            outer_except = code.funcstate.current_except
+            # Currently points to self, but the ExceptClauseNode would also be ok. Change if needed.
+            code.funcstate.current_except = self
             for except_clause in self.except_clauses:
                 except_clause.generate_handling_code(code, except_end_label)
+            code.funcstate.current_except = outer_except
+
             if not self.has_default_clause:
                 code.put_goto(except_error_label)
 
@@ -6813,7 +6823,6 @@ class TryExceptStatNode(StatNode):
                 code.put_label(exit_label)
                 code.mark_pos(self.pos, trace=False)
                 if can_raise:
-                    code.putln("__Pyx_PyThreadState_assign")  # re-assign in case a generator yielded
                     restore_saved_exception()
                 code.put_goto(old_label)
 
@@ -6822,7 +6831,6 @@ class TryExceptStatNode(StatNode):
                 code.put_goto(try_end_label)
             code.put_label(except_end_label)
             if can_raise:
-                code.putln("__Pyx_PyThreadState_assign")  # re-assign in case a generator yielded
                 restore_saved_exception()
         if code.label_used(try_end_label):
             code.put_label(try_end_label)
@@ -7086,6 +7094,7 @@ class TryFinallyStatNode(StatNode):
 
         if preserve_error:
             code.putln('/*exception exit:*/{')
+            if not self.in_generator:
                 code.putln("__Pyx_PyThreadState_declare")
             if self.is_try_finally_in_nogil:
                 code.declare_gilstate()
@@ -7148,7 +7157,6 @@ class TryFinallyStatNode(StatNode):
             if old_label == return_label:
                 # return actually raises an (uncatchable) exception in generators that we must preserve
                 if self.in_generator:
-                    code.putln("__Pyx_PyThreadState_declare")
                     exc_vars = tuple([
                         code.funcstate.allocate_temp(py_object_type, manage_ref=False)
                         for _ in range(6)])
@@ -7229,8 +7237,6 @@ class TryFinallyStatNode(StatNode):
 
         if self.is_try_finally_in_nogil:
             code.put_ensure_gil(declare_gilstate=False)
-        if self.in_generator:
-            code.putln("__Pyx_PyThreadState_assign")  # re-assign in case a generator yielded
 
         # not using preprocessor here to avoid warnings about
         # unused utility functions and/or temps
@@ -7257,8 +7263,6 @@ class TryFinallyStatNode(StatNode):
         code.globalstate.use_utility_code(reset_exception_utility_code)
         if self.is_try_finally_in_nogil:
             code.put_ensure_gil(declare_gilstate=False)
-        if self.in_generator:
-            code.putln("__Pyx_PyThreadState_assign")  # re-assign in case a generator yielded
 
         # not using preprocessor here to avoid warnings about
         # unused utility functions and/or temps

Cython/Compiler/ParseTreeTransforms.pxd

@@ -52,6 +52,7 @@ cdef class YieldNodeCollector(TreeVisitor):
     cdef public list yields
     cdef public list returns
     cdef public list finallys
+    cdef public list excepts
     cdef public bint has_return_value
     cdef public bint has_yield
     cdef public bint has_await

Cython/Compiler/ParseTreeTransforms.py

@@ -2470,6 +2470,7 @@ class YieldNodeCollector(TreeVisitor):
         self.yields = []
         self.returns = []
         self.finallys = []
+        self.excepts = []
         self.has_return_value = False
         self.has_yield = False
         self.has_await = False
@@ -2497,6 +2498,10 @@ class YieldNodeCollector(TreeVisitor):
         self.visitchildren(node)
         self.finallys.append(node)
 
+    def visit_TryExceptStatNode(self, node):
+        self.visitchildren(node)
+        self.excepts.append(node)
+
     def visit_ClassDefNode(self, node):
         pass
 
@@ -2547,7 +2552,7 @@ class MarkClosureVisitor(CythonTransform):
 
         for i, yield_expr in enumerate(collector.yields, 1):
             yield_expr.label_num = i
-        for retnode in collector.returns + collector.finallys:
+        for retnode in collector.returns + collector.finallys + collector.excepts:
             retnode.in_generator = True
 
         gbody = Nodes.GeneratorBodyDefNode(

Cython/Utility/Coroutine.c

@@ -389,6 +389,18 @@ static PyObject *__Pyx_Coroutine_Send(PyObject *self, PyObject *value); /*proto*
 static PyObject *__Pyx_Coroutine_Close(PyObject *self); /*proto*/
 static PyObject *__Pyx_Coroutine_Throw(PyObject *gen, PyObject *args); /*proto*/
 
+// macros for exception state swapping instead of inline functions to make use of the local thread state context
+#define __Pyx_Coroutine_SwapException(self) { \
+    __Pyx_ExceptionSwap(&(self)->exc_type, &(self)->exc_value, &(self)->exc_traceback); \
+    __Pyx_Coroutine_ResetFrameBackpointer(self); \
+    }
+#define __Pyx_Coroutine_ResetAndClearException(self) { \
+    __Pyx_ExceptionReset((self)->exc_type, (self)->exc_value, (self)->exc_traceback); \
+    __Pyx_Coroutine_ResetFrameBackpointer(self); \
+    (self)->exc_type = (self)->exc_value = (self)->exc_traceback = NULL; \
+    }
+static CYTHON_INLINE void __Pyx_Coroutine_ResetFrameBackpointer(__pyx_CoroutineObject *self);
+
 #if 1 || PY_VERSION_HEX < 0x030300B0
 static int __Pyx_PyGen_FetchStopIterationValue(PyObject **pvalue); /*proto*/
 #else
@@ -621,7 +633,6 @@ static
 PyObject *__Pyx_Coroutine_SendEx(__pyx_CoroutineObject *self, PyObject *value, int closing) {
     __Pyx_PyThreadState_declare
     PyObject *retval;
-    int restore_exc;
 
     assert(!self->is_running);
 
@@ -656,41 +667,32 @@ PyObject *__Pyx_Coroutine_SendEx(__pyx_CoroutineObject *self, PyObject *value, i
         __Pyx_ExceptionSwap(&self->exc_type, &self->exc_value,
                             &self->exc_traceback);
         // self->exc_* now holds the exception state of the caller
-        restore_exc = 1;
     } else {
         // save away the exception state of the caller
         __Pyx_Coroutine_ExceptionClear(self);
-        __Pyx_ExceptionSave(&self->exc_type, &self->exc_value,
-                            &self->exc_traceback);
-        restore_exc = 0;
+        __Pyx_ExceptionSave(&self->exc_type, &self->exc_value, &self->exc_traceback);
     }
 
     self->is_running = 1;
     retval = self->body((PyObject *) self, value);
     self->is_running = 0;
 
-    if (retval) {
-        // restore exception state of caller and save the state of the coroutine
-        __Pyx_ExceptionSwap(&self->exc_type, &self->exc_value,
-                            &self->exc_traceback);
-    }
+    return retval;
+}
+
+static CYTHON_INLINE void __Pyx_Coroutine_ResetFrameBackpointer(__pyx_CoroutineObject *self) {
 #if CYTHON_COMPILING_IN_PYPY || CYTHON_COMPILING_IN_PYSTON
     // FIXME: what to do in PyPy?
 #else
     // Don't keep the reference to f_back any longer than necessary.  It
     // may keep a chain of frames alive or it could create a reference
     // cycle.
-    if (restore_exc && self->exc_traceback) {
+    if (self->exc_traceback) {
         PyTracebackObject *tb = (PyTracebackObject *) self->exc_traceback;
         PyFrameObject *f = tb->tb_frame;
         Py_CLEAR(f->f_back);
     }
 #endif
-    if (!restore_exc) {
-        __Pyx_Coroutine_ExceptionClear(self);
-    }
-
-    return retval;
 }
 
 static CYTHON_INLINE