Correctly link calls when in-between calls are filtered.

pycallgraph/tracer.py

@@ -7,6 +7,7 @@
 from distutils import sysconfig
 from collections import defaultdict
 from threading import Thread
+
 try:
     from Queue import Queue, Empty
 except ImportError:
@@ -16,7 +17,6 @@
 
 
 class SyncronousTracer(object):
-
     def __init__(self, outputs, config):
         self.processor = TraceProcessor(outputs, config)
         self.config = config
@@ -41,7 +41,6 @@ def done(self):
 
 
 class AsyncronousTracer(SyncronousTracer):
-
     def start(self):
         self.processor.start()
         SyncronousTracer.start(self)
@@ -81,6 +80,10 @@ def init_trace_data(self):
         # Current call stack
         self.call_stack = ['__main__']
 
+        # Steps to the parent function
+        self.call_stack_back_steps = []
+        self.call_stack_back_steps.append(1)
+
         # Counters for each function
         self.func_count = defaultdict(int)
         self.func_count_max = 0
@@ -158,6 +161,7 @@ def process(self, frame, event, arg, memory=None):
                 )
 
         if event == 'call':
+
             keep = True
             code = frame.f_code
 
@@ -210,10 +214,13 @@ def process(self, frame, event, arg, memory=None):
             if keep:
 
                 if self.call_stack:
-                    src_func = self.call_stack[-1]
+                    src_func = self.call_stack[-self.call_stack_back_steps[-1]]
                 else:
                     src_func = None
 
+                # Start counting forward steps
+                self.call_stack_back_steps.append(1)
+
                 self.call_dict[src_func][full_name] += 1
 
                 self.func_count[full_name] += 1
@@ -229,6 +236,8 @@ def process(self, frame, event, arg, memory=None):
                     self.call_stack_memory_out.append([full_name, memory])
 
             else:
+                # Increment back steps
+                self.call_stack_back_steps[-1] += 1
                 self.call_stack.append('')
                 self.call_stack_timer.append(None)
 
@@ -239,6 +248,12 @@ def process(self, frame, event, arg, memory=None):
             if self.call_stack:
                 full_name = self.call_stack.pop(-1)
 
+                # Reset back steps
+                self.call_stack_back_steps[-1] -= 1
+                # If there are no steps to go back we remove the element
+                if self.call_stack_back_steps[-1] == 0:
+                    self.call_stack_back_steps.pop(-1)
+
                 if self.call_stack_timer:
                     start_time = self.call_stack_timer.pop(-1)
                 else:
@@ -372,4 +387,5 @@ def wrapper(*rest):
 
     return wrapper
 
+
 inspect.getmodule = simple_memoize(inspect.getmodule)