Implement line tracing for generators and coroutines.

Closes #1949.

CHANGES.rst

@@ -34,6 +34,9 @@ Features added
 Bugs fixed
 ----------
 
+* Line tracing did not include generators and coroutines.
+  (Github issue #1949)
+
 * C++ declarations for ``unordered_map`` were corrected.
   Patch by Michael Schatzow.  (Github issue #1484)
 

Cython/Compiler/Nodes.py

@@ -4032,9 +4032,10 @@ class GeneratorDefNode(DefNode):
 
         code.putln('{')
         code.putln('__pyx_CoroutineObject *gen = __Pyx_%s_New('
-                   '(__pyx_coroutine_body_t) %s, (PyObject *) %s, %s, %s, %s); %s' % (
+                   '(__pyx_coroutine_body_t) %s, %s, (PyObject *) %s, %s, %s, %s); %s' % (
                        self.gen_type_name,
-                       body_cname, Naming.cur_scope_cname, name, qualname, module_name,
+                       body_cname, self.code_object.calculate_result_code(code) if self.code_object else 'NULL',
+                       Naming.cur_scope_cname, name, qualname, module_name,
                        code.error_goto_if_null('gen', self.pos)))
         code.put_decref(Naming.cur_scope_cname, py_object_type)
         if self.requires_classobj:
@@ -4129,6 +4130,9 @@ class GeneratorBodyDefNode(DefNode):
         linetrace = code.globalstate.directives['linetrace']
         if profile or linetrace:
             tempvardecl_code.put_trace_declarations()
+            code.funcstate.can_trace = True
+            code_object = self.code_object.calculate_result_code(code) if self.code_object else None
+            code.put_trace_frame_init(code_object)
 
         # ----- Resume switch point.
         code.funcstate.init_closure_temps(lenv.scope_class.type.scope)
@@ -4167,6 +4171,9 @@ class GeneratorBodyDefNode(DefNode):
             # FIXME: this silences a potential "unused" warning => try to avoid unused closures in more cases
             code.putln("CYTHON_MAYBE_UNUSED_VAR(%s);" % Naming.cur_scope_cname)
 
+        if profile or linetrace:
+            code.funcstate.can_trace = False
+
         code.mark_pos(self.pos)
         code.putln("")
         code.putln("/* function exit code */")

Cython/Compiler/ParseTreeTransforms.py

@@ -1940,6 +1940,8 @@ if VALUE is not None:
             binding = self.current_directives.get('binding')
             rhs = ExprNodes.PyCFunctionNode.from_defnode(node, binding)
             node.code_object = rhs.code_object
+            if node.is_generator:
+                node.gbody.code_object = node.code_object
 
         if env.is_py_class_scope:
             rhs.binding = True

Cython/Utility/AsyncGen.c

@@ -34,7 +34,7 @@ static PyObject *__Pyx__PyAsyncGenValueWrapperNew(PyObject *val);
 
 
 static __pyx_CoroutineObject *__Pyx_AsyncGen_New(
-            __pyx_coroutine_body_t body, PyObject *closure,
+            __pyx_coroutine_body_t body, PyObject *code, PyObject *closure,
             PyObject *name, PyObject *qualname, PyObject *module_name) {
     __pyx_PyAsyncGenObject *gen = PyObject_GC_New(__pyx_PyAsyncGenObject, __pyx_AsyncGenType);
     if (unlikely(!gen))
@@ -42,7 +42,7 @@ static __pyx_CoroutineObject *__Pyx_AsyncGen_New(
     gen->ag_finalizer = NULL;
     gen->ag_closed = 0;
     gen->ag_hooks_inited = 0;
-    return __Pyx__Coroutine_NewInit((__pyx_CoroutineObject*)gen, body, closure, name, qualname, module_name);
+    return __Pyx__Coroutine_NewInit((__pyx_CoroutineObject*)gen, body, code, closure, name, qualname, module_name);
 }
 
 static int __pyx_AsyncGen_init(void);

Cython/Utility/Coroutine.c

@@ -382,17 +382,18 @@ typedef struct {
     PyObject *gi_name;
     PyObject *gi_qualname;
     PyObject *gi_modulename;
+    PyObject *gi_code;
     int resume_label;
     // using T_BOOL for property below requires char value
     char is_running;
 } __pyx_CoroutineObject;
 
 static __pyx_CoroutineObject *__Pyx__Coroutine_New(
-    PyTypeObject *type, __pyx_coroutine_body_t body, PyObject *closure,
+    PyTypeObject *type, __pyx_coroutine_body_t body, PyObject *code, PyObject *closure,
     PyObject *name, PyObject *qualname, PyObject *module_name); /*proto*/
 
 static __pyx_CoroutineObject *__Pyx__Coroutine_NewInit(
-            __pyx_CoroutineObject *gen, __pyx_coroutine_body_t body, PyObject *closure,
+            __pyx_CoroutineObject *gen, __pyx_coroutine_body_t body, PyObject *code, PyObject *closure,
             PyObject *name, PyObject *qualname, PyObject *module_name); /*proto*/
 
 static int __Pyx_Coroutine_clear(PyObject *self); /*proto*/
@@ -429,8 +430,8 @@ static PyTypeObject *__pyx_CoroutineAwaitType = 0;
 #define __Pyx_Coroutine_CheckExact(obj) (Py_TYPE(obj) == __pyx_CoroutineType)
 #define __Pyx_CoroutineAwait_CheckExact(obj) (Py_TYPE(obj) == __pyx_CoroutineAwaitType)
 
-#define __Pyx_Coroutine_New(body, closure, name, qualname, module_name)  \
-    __Pyx__Coroutine_New(__pyx_CoroutineType, body, closure, name, qualname, module_name)
+#define __Pyx_Coroutine_New(body, code, closure, name, qualname, module_name)  \
+    __Pyx__Coroutine_New(__pyx_CoroutineType, body, code, closure, name, qualname, module_name)
 
 static int __pyx_Coroutine_init(void); /*proto*/
 static PyObject *__Pyx__Coroutine_await(PyObject *coroutine); /*proto*/
@@ -450,8 +451,8 @@ static PyObject *__Pyx_CoroutineAwait_Throw(__pyx_CoroutineAwaitObject *self, Py
 static PyTypeObject *__pyx_GeneratorType = 0;
 #define __Pyx_Generator_CheckExact(obj) (Py_TYPE(obj) == __pyx_GeneratorType)
 
-#define __Pyx_Generator_New(body, closure, name, qualname, module_name)  \
-    __Pyx__Coroutine_New(__pyx_GeneratorType, body, closure, name, qualname, module_name)
+#define __Pyx_Generator_New(body, code, closure, name, qualname, module_name)  \
+    __Pyx__Coroutine_New(__pyx_GeneratorType, body, code, closure, name, qualname, module_name)
 
 static PyObject *__Pyx_Generator_Next(PyObject *self);
 static int __pyx_Generator_init(void); /*proto*/
@@ -1056,6 +1057,7 @@ static int __Pyx_Coroutine_clear(PyObject *self) {
         Py_CLEAR(((__pyx_PyAsyncGenObject*)gen)->ag_finalizer);
     }
 #endif
+    Py_CLEAR(gen->gi_code);
     Py_CLEAR(gen->gi_name);
     Py_CLEAR(gen->gi_qualname);
     Py_CLEAR(gen->gi_modulename);
@@ -1287,16 +1289,16 @@ __Pyx_Coroutine_set_qualname(__pyx_CoroutineObject *self, PyObject *value)
 }
 
 static __pyx_CoroutineObject *__Pyx__Coroutine_New(
-            PyTypeObject* type, __pyx_coroutine_body_t body, PyObject *closure,
+            PyTypeObject* type, __pyx_coroutine_body_t body, PyObject *code, PyObject *closure,
             PyObject *name, PyObject *qualname, PyObject *module_name) {
     __pyx_CoroutineObject *gen = PyObject_GC_New(__pyx_CoroutineObject, type);
     if (unlikely(!gen))
         return NULL;
-    return __Pyx__Coroutine_NewInit(gen, body, closure, name, qualname, module_name);
+    return __Pyx__Coroutine_NewInit(gen, body, code, closure, name, qualname, module_name);
 }
 
 static __pyx_CoroutineObject *__Pyx__Coroutine_NewInit(
-            __pyx_CoroutineObject *gen, __pyx_coroutine_body_t body, PyObject *closure,
+            __pyx_CoroutineObject *gen, __pyx_coroutine_body_t body, PyObject *code, PyObject *closure,
             PyObject *name, PyObject *qualname, PyObject *module_name) {
     gen->body = body;
     gen->closure = closure;
@@ -1315,6 +1317,8 @@ static __pyx_CoroutineObject *__Pyx__Coroutine_NewInit(
     gen->gi_name = name;
     Py_XINCREF(module_name);
     gen->gi_modulename = module_name;
+    Py_XINCREF(code);
+    gen->gi_code = code;
 
     PyObject_GC_Track(gen);
     return gen;
@@ -1485,6 +1489,7 @@ static PyMemberDef __pyx_Coroutine_memberlist[] = {
     {(char *) "cr_running", T_BOOL, offsetof(__pyx_CoroutineObject, is_running), READONLY, NULL},
     {(char*) "cr_await", T_OBJECT, offsetof(__pyx_CoroutineObject, yieldfrom), READONLY,
      (char*) PyDoc_STR("object being awaited, or None")},
+    {(char*) "cr_code", T_OBJECT, offsetof(__pyx_CoroutineObject, gi_code), READONLY, NULL},
     {(char *) "__module__", T_OBJECT, offsetof(__pyx_CoroutineObject, gi_modulename), PY_WRITE_RESTRICTED, 0},
     {0, 0, 0, 0, 0}
 };
@@ -1606,6 +1611,7 @@ static PyMemberDef __pyx_Generator_memberlist[] = {
     {(char *) "gi_running", T_BOOL, offsetof(__pyx_CoroutineObject, is_running), READONLY, NULL},
     {(char*) "gi_yieldfrom", T_OBJECT, offsetof(__pyx_CoroutineObject, yieldfrom), READONLY,
      (char*) PyDoc_STR("object being iterated by 'yield from', or None")},
+    {(char*) "gi_code", T_OBJECT, offsetof(__pyx_CoroutineObject, gi_code), READONLY, NULL},
     {0, 0, 0, 0, 0}
 };
 

tests/run/line_profile_test.srctree

@@ -49,6 +49,18 @@ profile = line_profiler.LineProfiler(func)
 profile.runcall(func, 19)
 assert_stats(profile, func.__name__)
 
+from collatz import run_generator, cy_generator
+func = cy_generator
+profile = line_profiler.LineProfiler(func)
+profile.runcall(run_generator, 19)
+assert_stats(profile, func.__name__)
+
+from collatz import run_coro, cy_coro
+func = cy_coro
+profile = line_profiler.LineProfiler(func)
+profile.runcall(run_coro, 19)
+assert_stats(profile, func.__name__)
+
 from collatz import PyClass
 obj = PyClass()
 func = obj.py_pymethod
@@ -92,6 +104,38 @@ cpdef cp_collatz(n):
             n = 3*n+1
 
 
+@cython.binding(True)
+def cy_generator(int n):
+   x = 1
+   for i in range(n):
+       yield x + 2
+
+
+@cython.binding(True)
+def run_generator(n):
+    assert len(list(cy_generator(n))) == n
+
+
+@cython.binding(True)
+async def cy_coro(int n):
+    while n > 1:
+        if n % 2 == 0:
+            n //= 2
+        else:
+            n = 3*n+1
+
+
+@cython.binding(True)
+def run_coro(n):
+    coro = cy_coro(n)
+    try:
+        coro.send(None)
+    except StopIteration:
+        assert True
+    else:
+        assert False, "Coroutine did not raise"
+
+
 @cython.binding(True)
 class PyClass(object):
     def py_pymethod(self):