diff --git a/pycallgraph/tracer.py b/pycallgraph/tracer.py index ffcab158..21b2efa6 100644 --- a/pycallgraph/tracer.py +++ b/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)