Repository: cython/cython

Issue:
Profiling doesn't work for generators or generator expressions
For generators, the profiler 'call' hook is only called once on creation and the 'return' hook is never called. This creates an invalid call graph in pstats (as the profiler treats the generator as never exited) and skews the reported results.

In CPython, generators call the 'call' and 'return' hook for each iteration, and once again for the final StopIteration.


Discussion:

Resolving patch:
diff --git a/Cython/Compiler/ExprNodes.py b/Cython/Compiler/ExprNodes.py
--- a/Cython/Compiler/ExprNodes.py
+++ b/Cython/Compiler/ExprNodes.py
@@ -9287,7 +9287,13 @@ def generate_yield_code(self, code):
             code.putln('%s->%s = %s;' % (Naming.cur_scope_cname, save_cname, cname))
 
         code.put_xgiveref(Naming.retval_cname)
+        profile = code.globalstate.directives['profile']
+        linetrace = code.globalstate.directives['linetrace']
+        if profile or linetrace:
+            code.put_trace_return(Naming.retval_cname,
+                                  nogil=not code.funcstate.gil_owned)
         code.put_finish_refcount_context()
+
         code.putln("/* return from generator, yielding value */")
         code.putln("%s->resume_label = %d;" % (
             Naming.generator_cname, label_num))
diff --git a/Cython/Compiler/Nodes.py b/Cython/Compiler/Nodes.py
--- a/Cython/Compiler/Nodes.py
+++ b/Cython/Compiler/Nodes.py
@@ -1805,9 +1805,11 @@ def generate_function_definitions(self, env, code):
             code.declare_gilstate()
 
         if profile or linetrace:
-            tempvardecl_code.put_trace_declarations()
-            code_object = self.code_object.calculate_result_code(code) if self.code_object else None
-            code.put_trace_frame_init(code_object)
+            if not self.is_generator:
+                # generators are traced when iterated, not at creation
+                tempvardecl_code.put_trace_declarations()
+                code_object = self.code_object.calculate_result_code(code) if self.code_object else None
+                code.put_trace_frame_init(code_object)
 
         # ----- set up refnanny
         if use_refnanny:
@@ -1862,12 +1864,14 @@ def generate_function_definitions(self, env, code):
         if profile or linetrace:
             # this looks a bit late, but if we don't get here due to a
             # fatal error before hand, it's not really worth tracing
-            if self.is_wrapper:
-                trace_name = self.entry.name + " (wrapper)"
-            else:
-                trace_name = self.entry.name
-            code.put_trace_call(
-                trace_name, self.pos, nogil=not code.funcstate.gil_owned)
+            if not self.is_generator:
+                # generators are traced when iterated, not at creation
+                if self.is_wrapper:
+                    trace_name = self.entry.name + " (wrapper)"
+                else:
+                    trace_name = self.entry.name
+                code.put_trace_call(
+                    trace_name, self.pos, nogil=not code.funcstate.gil_owned)
             code.funcstate.can_trace = True
         # ----- Fetch arguments
         self.generate_argument_parsing_code(env, code)
@@ -2064,12 +2068,14 @@ def generate_function_definitions(self, env, code):
 
         if profile or linetrace:
             code.funcstate.can_trace = False
-            if self.return_type.is_pyobject:
-                code.put_trace_return(
-                    Naming.retval_cname, nogil=not code.funcstate.gil_owned)
-            else:
-                code.put_trace_return(
-                    "Py_None", nogil=not code.funcstate.gil_owned)
+            if not self.is_generator:
+                # generators are traced when iterated, not at creation
+                if self.return_type.is_pyobject:
+                    code.put_trace_return(
+                        Naming.retval_cname, nogil=not code.funcstate.gil_owned)
+                else:
+                    code.put_trace_return(
+                        "Py_None", nogil=not code.funcstate.gil_owned)
 
         if not lenv.nogil:
             # GIL holding function
@@ -4041,6 +4047,10 @@ def generate_function_definitions(self, env, code):
         tempvardecl_code = code.insertion_point()
         code.put_declare_refcount_context()
         code.put_setup_refcount_context(self.entry.name)
+        profile = code.globalstate.directives['profile']
+        linetrace = code.globalstate.directives['linetrace']
+        if profile or linetrace:
+            code.put_trace_declarations()
 
         # ----- Resume switch point.
         code.funcstate.init_closure_temps(lenv.scope_class.type.scope)
@@ -4112,6 +4122,9 @@ def generate_function_definitions(self, env, code):
         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)
+        if profile or linetrace:
+            code.put_trace_return(Naming.retval_cname,
+                                  nogil=not code.funcstate.gil_owned)
         code.put_finish_refcount_context()
         code.putln("return %s;" % Naming.retval_cname)
         code.putln("}")
@@ -4119,13 +4132,20 @@ def generate_function_definitions(self, env, code):
         # ----- Go back and insert temp variable declarations
         tempvardecl_code.put_temp_declarations(code.funcstate)
         # ----- Generator resume code
+        if profile or linetrace:
+            resume_code.put_trace_call(self.entry.qualified_name, self.pos,
+                                       nogil=not code.funcstate.gil_owned)
         resume_code.putln("switch (%s->resume_label) {" % (
                        Naming.generator_cname))
+
         resume_code.putln("case 0: goto %s;" % first_run_label)
 
         for i, label in code.yield_labels:
             resume_code.putln("case %d: goto %s;" % (i, label))
         resume_code.putln("default: /* CPython raises the right error here */")
+        if profile or linetrace:
+            resume_code.put_trace_return("Py_None",
+                                         nogil=not code.funcstate.gil_owned)
         resume_code.put_finish_refcount_context()
         resume_code.putln("return NULL;")
         resume_code.putln("}")


Test patch:
diff --git a/tests/run/pstats_profile_test.pyx b/tests/run/pstats_profile_test.pyx
--- a/tests/run/pstats_profile_test.pyx
+++ b/tests/run/pstats_profile_test.pyx
@@ -65,6 +65,58 @@ __doc__ = u"""
      'f_raise',
      'm_cdef', 'm_cpdef', 'm_cpdef (wrapper)', 'm_def',
      'withgil_prof']
+
+    >>> profile.runctx("test_generators()", locals(), globals(), statsfile)
+    >>> s = pstats.Stats(statsfile)
+    >>> short_stats = dict([(k[2], v[1]) for k,v in s.stats.items()])
+    >>> short_stats['generator']
+    3
+
+    >>> short_stats['generator_exception']
+    2
+
+    >>> short_stats['genexpr']
+    11
+
+    >>> sorted(callees(s, 'test_generators'))
+    ['call_generator', 'call_generator_exception', 'generator_expr']
+
+    >>> list(callees(s, 'call_generator'))
+    ['generator']
+
+    >>> list(callees(s, 'generator'))
+    []
+
+    >>> list(callees(s, 'generator_exception'))
+    []
+
+    >>> list(callees(s, 'generator_expr'))
+    ['genexpr']
+
+    >>> list(callees(s, 'genexpr'))
+    []
+
+    >>> def python_generator():
+    ...   yield 1
+    ...   yield 2
+    >>> def call_python_generator():
+    ...   list(python_generator())
+
+    >>> profile.runctx("call_python_generator()", locals(), globals(), statsfile)
+    >>> python_stats = pstats.Stats(statsfile)
+    >>> python_stats_dict = dict([(k[2], v[1]) for k,v in python_stats.stats.items()])
+
+    >>> profile.runctx("call_generator()", locals(), globals(), statsfile)
+    >>> cython_stats = pstats.Stats(statsfile)
+    >>> cython_stats_dict = dict([(k[2], v[1]) for k,v in cython_stats.stats.items()])
+
+    >>> python_stats_dict['python_generator'] == cython_stats_dict['generator']
+    True
+
+    >>> try:
+    ...    os.unlink(statsfile)
+    ... except:
+    ...    pass
 """
 
 cimport cython
@@ -147,3 +199,29 @@ cdef class A(object):
         return a
     cdef m_cdef(self, long a):
         return a
+
+def test_generators():
+    call_generator()
+    call_generator_exception()
+    generator_expr()
+
+def call_generator():
+    list(generator())
+
+def generator():
+    yield 1
+    yield 2
+
+def call_generator_exception():
+    try:
+        list(generator_exception())
+    except ValueError:
+        pass
+
+def generator_exception():
+    yield 1
+    raise ValueError(2)
+
+def generator_expr():
+    e = (x for x in range(10))
+    return sum(e)