report variable changes from the last line in a with block

pysnooper/tracer.py

@@ -3,6 +3,7 @@
 
 import functools
 import inspect
+import dis
 import opcode
 import os
 import sys
@@ -19,6 +20,20 @@
     from io import open
 
 
+def _get_line_number(code, offset):
+    # Line number for a bytecode offset, from the code object's line table.
+    # frame.f_lineno at __enter__ is not reliable for the `with` line when the
+    # same statement runs again in a loop (older Pythons), but f_lasti is, so we
+    # resolve the line from it instead.
+    line_number = None
+    best = -1
+    for start, line in dis.findlinestarts(code):
+        if line is not None and best < start <= offset:
+            best = start
+            line_number = line
+    return line_number
+
+
 ipython_filename_pattern = re.compile('^<ipython-input-([0-9]+)-.*>$')
 ansible_filename_pattern = re.compile(r'^(.+\.zip)[/|\\](ansible[/|\\]modules[/|\\].+\.py)$')
 ipykernel_filename_pattern = re.compile(r'^/var/folders/.*/ipykernel_[0-9]+/[0-9]+.py$')
@@ -264,6 +279,8 @@ def __init__(self, output=None, watch=(), watch_explode=(), depth=1,
              for v in utils.ensure_tuple(watch_explode)
         ]
         self.frame_to_local_reprs = {}
+        self.frame_to_with_line = {}
+        self.frame_to_pending_flush = {}
         self.start_times = {}
         self.depth = depth
         self.prefix = prefix
@@ -372,6 +389,10 @@ def __enter__(self):
         if not self._is_internal_frame(calling_frame):
             calling_frame.f_trace = self.trace
             self.target_frames.add(calling_frame)
+            self.frame_to_with_line[calling_frame] = _get_line_number(
+                calling_frame.f_code, calling_frame.f_lasti
+            )
+            self.frame_to_pending_flush[calling_frame] = False
 
         stack = self.thread_local.__dict__.setdefault(
             'original_trace_functions', []
@@ -387,7 +408,32 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
         sys.settrace(stack.pop())
         calling_frame = inspect.currentframe().f_back
         self.target_frames.discard(calling_frame)
-        self.frame_to_local_reprs.pop(calling_frame, None)
+        start_time = self.start_times.pop(calling_frame)
+        # Variable changes are reported on the *next* trace event. The last line
+        # of a `with` block may have no next event before we get here (Python <
+        # 3.10, or a one-line body), so flush it now. We skip it when a later
+        # event already captured the final state: the block-exit line event on
+        # 3.10+ (which clears the pending flag), or the exception events on an
+        # exceptional exit. Otherwise every repr / watch / custom_repr would be
+        # evaluated twice, and on an exceptional exit a failing callback could
+        # even replace the user's exception. The frame is not in
+        # frame_to_local_reprs when the body produced no trace event at all (a
+        # one-line body), which still needs a flush. Cleanup runs in finally so
+        # a failing callback can't strand the per-frame state.
+        try:
+            # Only real `with` blocks (not the internal frame of the decorator
+            # form) are registered in frame_to_with_line.
+            is_with_block = calling_frame in self.frame_to_with_line
+            pending = self.frame_to_pending_flush.get(calling_frame, False)
+            never_reported = calling_frame not in self.frame_to_local_reprs
+            if exc_type is None and is_with_block and (pending or never_reported):
+                self._report_variable_changes(
+                    calling_frame, ' ' * 4 * thread_global.depth
+                )
+        finally:
+            self.frame_to_local_reprs.pop(calling_frame, None)
+            self.frame_to_with_line.pop(calling_frame, None)
+            self.frame_to_pending_flush.pop(calling_frame, None)
 
         ### Writing elapsed time: #############################################
         #                                                                     #
@@ -396,7 +442,6 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
         _STYLE_NORMAL = self._STYLE_NORMAL
         _STYLE_RESET_ALL = self._STYLE_RESET_ALL
 
-        start_time = self.start_times.pop(calling_frame)
         duration = datetime_module.datetime.now() - start_time
         elapsed_time_string = pycompat.timedelta_format(duration)
         indent = ' ' * 4 * (thread_global.depth + 1)
@@ -408,6 +453,32 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
         #                                                                     #
         ### Finished writing elapsed time. ####################################
 
+    def _report_variable_changes(self, frame, indent, is_call=False):
+        _FOREGROUND_GREEN = self._FOREGROUND_GREEN
+        _STYLE_DIM = self._STYLE_DIM
+        _STYLE_NORMAL = self._STYLE_NORMAL
+        _STYLE_RESET_ALL = self._STYLE_RESET_ALL
+
+        old_local_reprs = self.frame_to_local_reprs.get(frame, {})
+        self.frame_to_local_reprs[frame] = local_reprs = \
+                                       get_local_reprs(frame,
+                                                       watch=self.watch, custom_repr=self.custom_repr,
+                                                       max_length=self.max_variable_length,
+                                                       normalize=self.normalize,
+                                                       )
+
+        newish_string = 'Starting var:.. ' if is_call else 'New var:....... '
+
+        for name, value_repr in local_reprs.items():
+            if name not in old_local_reprs:
+                self.write('{indent}{_FOREGROUND_GREEN}{_STYLE_DIM}'
+                           '{newish_string}{_STYLE_NORMAL}{name} = '
+                           '{value_repr}{_STYLE_RESET_ALL}'.format(**locals()))
+            elif old_local_reprs[name] != value_repr:
+                self.write('{indent}{_FOREGROUND_GREEN}{_STYLE_DIM}'
+                           'Modified var:.. {_STYLE_NORMAL}{name} = '
+                           '{value_repr}{_STYLE_RESET_ALL}'.format(**locals()))
+
     def _is_internal_frame(self, frame):
         return frame.f_code.co_filename == Tracer.__enter__.__code__.co_filename
 
@@ -504,27 +575,17 @@ def trace(self, frame, event, arg):
 
         ### Reporting newish and modified variables: ##########################
         #                                                                     #
-        old_local_reprs = self.frame_to_local_reprs.get(frame, {})
-        self.frame_to_local_reprs[frame] = local_reprs = \
-                                       get_local_reprs(frame,
-                                                       watch=self.watch, custom_repr=self.custom_repr,
-                                                       max_length=self.max_variable_length,
-                                                       normalize=self.normalize,
-                                                       )
-
-        newish_string = ('Starting var:.. ' if event == 'call' else
-                                                            'New var:....... ')
-
-        for name, value_repr in local_reprs.items():
-            if name not in old_local_reprs:
-                self.write('{indent}{_FOREGROUND_GREEN}{_STYLE_DIM}'
-                           '{newish_string}{_STYLE_NORMAL}{name} = '
-                           '{value_repr}{_STYLE_RESET_ALL}'.format(**locals()))
-            elif old_local_reprs[name] != value_repr:
-                self.write('{indent}{_FOREGROUND_GREEN}{_STYLE_DIM}'
-                           'Modified var:.. {_STYLE_NORMAL}{name} = '
-                           '{value_repr}{_STYLE_RESET_ALL}'.format(**locals()))
-
+        self._report_variable_changes(frame, indent, is_call=(event == 'call'))
+        # Changes for a line are reported on the *next* event, so a body line
+        # leaves them pending. The block-exit line event (back on the `with`
+        # line, Python 3.10+) and non-line events like 'exception'/'return' have
+        # already captured the final state, so they clear the flag. __exit__
+        # uses it to flush the last line without re-reporting an already flushed
+        # one (which would re-run every repr / watch / custom_repr).
+        if frame in self.frame_to_with_line:
+            self.frame_to_pending_flush[frame] = (
+                event == 'line' and line_no != self.frame_to_with_line[frame]
+            )
         #                                                                     #
         ### Finished newish and modified variables. ###########################
 

tests/test_pysnooper.py

@@ -1190,7 +1190,7 @@ def f1(x1):
             LineEntry(),
             ReturnEntry(),
             ReturnValueEntry('20'),
-            VariableEntry(min_python_version=(3, 10)),
+            VariableEntry(),
             LineEntry(source_regex="with pysnooper.snoop.*", min_python_version=(3, 10)),
             ElapsedTimeEntry(),
         ),
@@ -1259,14 +1259,144 @@ def f1(a):
             ReturnValueEntry(),
             ReturnEntry(),
             ReturnValueEntry(),
-            VariableEntry(min_python_version=(3, 10)),
+            VariableEntry(),
             LineEntry(source_regex="with pysnooper.snoop.*", min_python_version=(3, 10)),
             ElapsedTimeEntry(),
         ),
         normalize=normalize,
     )
 
 
+def test_with_block_reports_last_line_variables():
+    # Regression for #237: the last line of a `with` block is not followed by
+    # another trace event on Python < 3.10, so a variable created or modified
+    # there used to be dropped. Make sure a final line that both modifies an
+    # existing variable and creates a new one is reported, exactly once, on
+    # every Python version.
+    string_io = io.StringIO()
+
+    def f():
+        with pysnooper.snoop(string_io, color=False, normalize=True):
+            x = 1
+            y = 2
+            x = 3; z = 4
+
+    f()
+    assert_output(
+        string_io.getvalue(),
+        (
+            SourcePathEntry(),
+            VariableEntry(),  # the output stream local
+            LineEntry('x = 1'),
+            VariableEntry('x', '1', stage='new'),
+            LineEntry('y = 2'),
+            VariableEntry('y', '2', stage='new'),
+            LineEntry('x = 3; z = 4'),
+            VariableEntry('x', '3', stage='modified'),
+            VariableEntry('z', '4', stage='new'),
+            LineEntry(source_regex="with pysnooper.snoop.*",
+                      min_python_version=(3, 10)),
+            ElapsedTimeEntry(),
+        ),
+        normalize=True,
+    )
+
+
+def test_with_block_one_line_body():
+    # Regression for #237: a one-line `with` body produces no trace event at
+    # all, so its variables have to be reported from __exit__.
+    string_io = io.StringIO()
+
+    def f():
+        with pysnooper.snoop(string_io, color=False, normalize=True): answer = 42
+
+    f()
+    assert_output(
+        string_io.getvalue(),
+        (
+            VariableEntry('answer', '42', stage='new'),
+            VariableEntry(),  # the output stream local
+            ElapsedTimeEntry(),
+        ),
+        normalize=True,
+    )
+
+
+def test_with_block_exit_does_not_reevaluate_reprs():
+    # Regression for #237: on Python 3.10+ leaving the block already reports
+    # the final variables via a line event, so __exit__ must not evaluate the
+    # reprs a second time (which would double custom_repr / watch calls).
+    string_io = io.StringIO()
+    call_count = [0]
+
+    def count_repr(x):
+        call_count[0] += 1
+        return 'LIST'
+
+    def f():
+        with pysnooper.snoop(
+            string_io, color=False,
+            custom_repr=((lambda x: isinstance(x, list), count_repr),),
+        ):
+            data = [1, 2, 3]
+
+    f()
+    assert call_count[0] == 1
+
+
+def test_with_block_repeated_in_loop():
+    # Regression for #237: a `with` block reused across loop iterations must
+    # report each iteration's final value. Older Pythons don't reliably report
+    # the `with` line at __enter__, so a mechanism keyed on it would report only
+    # the first iteration and leave the rest stale.
+    string_io = io.StringIO()
+
+    def f():
+        for i in range(3):
+            with pysnooper.snoop(string_io, color=False, normalize=True):
+                x = i
+
+    f()
+    output = string_io.getvalue()
+    assert 'New var:....... x = 0' in output
+    assert 'Modified var:.. x = 1' in output
+    assert 'Modified var:.. x = 2' in output
+
+
+def test_with_block_exceptional_exit():
+    # Regression for #237: when the block exits because of an exception, the
+    # trace events have already captured the final state, so __exit__ must not
+    # snapshot again (re-running callbacks, which could even replace the user's
+    # exception). The user's exception must propagate and per-frame state must
+    # be cleaned up.
+    string_io = io.StringIO()
+    tracer = pysnooper.snoop(string_io, color=False)
+
+    exit_snapshots = [0]
+    original = tracer._report_variable_changes
+
+    def spy(frame, indent, is_call=False):
+        if sys._getframe(1).f_code.co_name == '__exit__':
+            exit_snapshots[0] += 1
+        return original(frame, indent, is_call)
+
+    tracer._report_variable_changes = spy
+
+    def f():
+        with tracer:
+            data = [1, 2, 3]
+            raise ValueError("boom")
+
+    with pytest.raises(ValueError):
+        f()
+
+    assert exit_snapshots[0] == 0
+    assert not tracer.frame_to_local_reprs
+    assert not tracer.frame_to_with_line
+    assert not tracer.frame_to_pending_flush
+    assert not tracer.start_times
+
+
 @pytest.mark.parametrize("normalize", (True, False))
 def test_var_order(normalize):
     string_io = io.StringIO()
@@ -1309,7 +1439,7 @@ def f(one, two, three, four):
             VariableEntry("seven", "7"),
             ReturnEntry(),
             ReturnValueEntry(),
-            VariableEntry("result", "None", min_python_version=(3, 10)),
+            VariableEntry("result", "None"),
             LineEntry(source_regex="with pysnooper.snoop.*", min_python_version=(3, 10)),
             ElapsedTimeEntry(),
         ),