diff --git a/pysnooper/tracer.py b/pysnooper/tracer.py index 5929c87..b1a027a 100644 --- a/pysnooper/tracer.py +++ b/pysnooper/tracer.py @@ -190,12 +190,13 @@ class Tracer: self.overwrite = overwrite self._did_overwrite = False assert self.depth >= 1 - self.target_code_object = None + self.target_codes = set() + self.target_frames = set() def __call__(self, function): - self.target_code_object = function.__code__ + self.target_codes.add(function.__code__) - def inner(function_, *args, **kwargs): + def inner(_, *args, **kwargs): with self: return function(*args, **kwargs) @@ -209,16 +210,24 @@ class Tracer: self._write(s) def __enter__(self): - if not self.target_code_object: - calling_frame = inspect.currentframe().f_back - self.target_code_object = calling_frame.f_code + calling_frame = inspect.currentframe().f_back + if not self._is_internal_frame(calling_frame): calling_frame.f_trace = self.trace + self.target_frames.add(calling_frame) self.original_trace_function = sys.gettrace() sys.settrace(self.trace) def __exit__(self, exc_type, exc_value, exc_traceback): sys.settrace(self.original_trace_function) + calling_frame = inspect.currentframe().f_back + self.target_frames.discard(calling_frame) + + def _should_trace_frame(self, frame): + return frame.f_code in self.target_codes or frame in self.target_frames + + def _is_internal_frame(self, frame): + return frame.f_code.co_filename == __file__ def trace(self, frame, event, arg): @@ -228,19 +237,21 @@ class Tracer: # or the user asked to go a few levels deeper and we're within that # number of levels deeper. - if frame.f_code is not self.target_code_object: + if not self._should_trace_frame(frame): if self.depth == 1: # We did the most common and quickest check above, because the # trace function runs so incredibly often, therefore it's # crucial to hyper-optimize it for the common case. return self.trace + elif self._is_internal_frame(frame): + return self.trace else: _frame_candidate = frame for i in range(1, self.depth): _frame_candidate = _frame_candidate.f_back if _frame_candidate is None: return self.trace - elif _frame_candidate.f_code is self.target_code_object: + elif self._should_trace_frame(_frame_candidate): indent = ' ' * 4 * i break else: diff --git a/tests/test_pysnooper.py b/tests/test_pysnooper.py index c9586b8..951131b 100644 --- a/tests/test_pysnooper.py +++ b/tests/test_pysnooper.py @@ -632,3 +632,180 @@ def test_needs_parentheses(): assert needs_parentheses('x * y') assert needs_parentheses('x and y') assert needs_parentheses('x if z else y') + + +def test_with_block(): + # Testing that a single Tracer can handle many mixed uses + snoop = pysnooper.snoop() + + def foo(x): + if x == 0: + bar1(x) + qux() + return + + with snoop: + # There should be line entries for these three lines, + # no line entries for anything else in this function, + # but calls to all bar functions should be traced + foo(x - 1) + bar2(x) + qux() + int(4) + bar3(9) + return x + + @snoop + def bar1(_x): + qux() + + @snoop + def bar2(_x): + qux() + + @snoop + def bar3(_x): + qux() + + def qux(): + return 9 # not traced, mustn't show up + + with sys_tools.OutputCapturer(stdout=False, + stderr=True) as output_capturer: + result = foo(2) + assert result == 2 + output = output_capturer.string_io.getvalue() + assert_output( + output, + ( + # In first with + VariableEntry('bar1'), + VariableEntry('bar2'), + VariableEntry('bar3'), + VariableEntry('foo'), + VariableEntry('qux'), + VariableEntry('snoop'), + VariableEntry('x', '2'), + LineEntry('foo(x - 1)'), + + # In with in recursive call + VariableEntry('bar1'), + VariableEntry('bar2'), + VariableEntry('bar3'), + VariableEntry('foo'), + VariableEntry('qux'), + VariableEntry('snoop'), + VariableEntry('x', '1'), + LineEntry('foo(x - 1)'), + + # Call to bar1 from if block outside with + VariableEntry('_x', '0'), + VariableEntry('qux'), + CallEntry('def bar1(_x):'), + LineEntry('qux()'), + ReturnEntry('qux()'), + ReturnValueEntry('None'), + + # In with in recursive call + LineEntry('bar2(x)'), + + # Call to bar2 from within with + VariableEntry('_x', '1'), + VariableEntry('qux'), + CallEntry('def bar2(_x):'), + LineEntry('qux()'), + ReturnEntry('qux()'), + ReturnValueEntry('None'), + + # In with in recursive call + LineEntry('qux()'), + + # Call to bar3 from after with + VariableEntry('_x', '9'), + VariableEntry('qux'), + CallEntry('def bar3(_x):'), + LineEntry('qux()'), + ReturnEntry('qux()'), + ReturnValueEntry('None'), + + # -- Similar to previous few sections, + # -- but from first call to foo + + # In with in first call + LineEntry('bar2(x)'), + + # Call to bar2 from within with + VariableEntry('_x', '2'), + VariableEntry('qux'), + CallEntry('def bar2(_x):'), + LineEntry('qux()'), + ReturnEntry('qux()'), + ReturnValueEntry('None'), + + # In with in first call + LineEntry('qux()'), + + # Call to bar3 from after with + VariableEntry('_x', '9'), + VariableEntry('qux'), + CallEntry('def bar3(_x):'), + LineEntry('qux()'), + ReturnEntry('qux()'), + ReturnValueEntry('None'), + ), + ) + + +def test_with_block_depth(): + string_io = io.StringIO() + + def f4(x4): + result4 = x4 * 2 + return result4 + + def f3(x3): + result3 = f4(x3) + return result3 + + def f2(x2): + result2 = f3(x2) + return result2 + + def f1(x1): + str(3) + with pysnooper.snoop(string_io, depth=3): + result1 = f2(x1) + return result1 + + result = f1(10) + assert result == 20 + output = string_io.getvalue() + assert_output( + output, + ( + VariableEntry(), + VariableEntry(), + VariableEntry(), + LineEntry('result1 = f2(x1)'), + + VariableEntry(), + VariableEntry(), + CallEntry('def f2(x2):'), + LineEntry(), + + VariableEntry(), + VariableEntry(), + CallEntry('def f3(x3):'), + LineEntry(), + + VariableEntry(), + LineEntry(), + ReturnEntry(), + ReturnValueEntry('20'), + + VariableEntry(), + LineEntry(), + ReturnEntry(), + ReturnValueEntry('20'), + ) + )