Merge pull request #1669 from shalabhc/generator-profiling-fix

Fix profiling for generators and generator expressions

Cython/Compiler/ExprNodes.py

@@ -9287,7 +9287,13 @@ class YieldExprNode(ExprNode):
             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))

Cython/Compiler/Nodes.py

@@ -1805,6 +1805,8 @@ class FuncDefNode(StatNode, BlockNode):
             code.declare_gilstate()
 
         if profile or linetrace:
+            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)
@@ -1862,6 +1864,8 @@ class FuncDefNode(StatNode, BlockNode):
         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 not self.is_generator:
+                # generators are traced when iterated, not at creation
                 if self.is_wrapper:
                     trace_name = self.entry.name + " (wrapper)"
                 else:
@@ -2064,6 +2068,8 @@ class FuncDefNode(StatNode, BlockNode):
 
         if profile or linetrace:
             code.funcstate.can_trace = False
+            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)
@@ -4041,6 +4047,10 @@ class GeneratorBodyDefNode(DefNode):
         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 @@ class GeneratorBodyDefNode(DefNode):
         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 @@ class GeneratorBodyDefNode(DefNode):
         # ----- 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("}")

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)