refactor: clarify the code that fixes with-statement exits

Author: nedbat

coverage/env.py

@@ -79,10 +79,66 @@ class PYBEHAVIOR:
     keep_constant_test = pep626
 
     # When leaving a with-block, do we visit the with-line again for the exit?
+    # For example, wwith.py:
+    #
+    #    with open("/tmp/test", "w") as f1:
+    #        a = 2
+    #        with open("/tmp/test2", "w") as f3:
+    #            print(4)
+    #
+    # % python3.9 -m trace -t wwith.py | grep wwith
+    #  --- modulename: wwith, funcname: <module>
+    # wwith.py(1): with open("/tmp/test", "w") as f1:
+    # wwith.py(2):     a = 2
+    # wwith.py(3):     with open("/tmp/test2", "w") as f3:
+    # wwith.py(4):         print(4)
+    #
+    # % python3.10 -m trace -t wwith.py | grep wwith
+    #  --- modulename: wwith, funcname: <module>
+    # wwith.py(1): with open("/tmp/test", "w") as f1:
+    # wwith.py(2):     a = 2
+    # wwith.py(3):     with open("/tmp/test2", "w") as f3:
+    # wwith.py(4):         print(4)
+    # wwith.py(3):     with open("/tmp/test2", "w") as f3:
+    # wwith.py(1): with open("/tmp/test", "w") as f1:
+    #
     exit_through_with = (PYVERSION >= (3, 10, 0, "beta"))
 
     # When leaving a with-block, do we visit the with-line exactly,
-    # or the inner-most context manager?
+    # or the context managers in inner-out order?
+    #
+    # mwith.py:
+    #    with (
+    #        open("/tmp/one", "w") as f2,
+    #        open("/tmp/two", "w") as f3,
+    #        open("/tmp/three", "w") as f4,
+    #    ):
+    #        print("hello 6")
+    #
+    # % python3.11 -m trace -t mwith.py | grep mwith
+    #  --- modulename: mwith, funcname: <module>
+    # mwith.py(2):     open("/tmp/one", "w") as f2,
+    # mwith.py(1): with (
+    # mwith.py(2):     open("/tmp/one", "w") as f2,
+    # mwith.py(3):     open("/tmp/two", "w") as f3,
+    # mwith.py(1): with (
+    # mwith.py(3):     open("/tmp/two", "w") as f3,
+    # mwith.py(4):     open("/tmp/three", "w") as f4,
+    # mwith.py(1): with (
+    # mwith.py(4):     open("/tmp/three", "w") as f4,
+    # mwith.py(6):     print("hello 6")
+    # mwith.py(1): with (
+    #
+    # % python3.12 -m trace -t mwith.py | grep mwith
+    #  --- modulename: mwith, funcname: <module>
+    # mwith.py(2):      open("/tmp/one", "w") as f2,
+    # mwith.py(3):      open("/tmp/two", "w") as f3,
+    # mwith.py(4):      open("/tmp/three", "w") as f4,
+    # mwith.py(6):      print("hello 6")
+    # mwith.py(4):      open("/tmp/three", "w") as f4,
+    # mwith.py(3):      open("/tmp/two", "w") as f3,
+    # mwith.py(2):      open("/tmp/one", "w") as f2,
+
     exit_with_through_ctxmgr = (PYVERSION >= (3, 12, 6))
 
     # Match-case construct.

coverage/parser.py

@@ -300,10 +300,14 @@ def _analyze_ast(self) -> None:
         assert self._ast_root is not None
         aaa = AstArcAnalyzer(self.filename, self._ast_root, self.raw_statements, self._multiline)
         aaa.analyze()
-        self._with_jump_fixers = aaa.with_jump_fixers()
+        arcs = aaa.arcs
+        if env.PYBEHAVIOR.exit_through_with:
+            self._with_jump_fixers = aaa.with_jump_fixers()
+            if self._with_jump_fixers:
+                arcs = self.fix_with_jumps(arcs)
 
         self._all_arcs = set()
-        for l1, l2 in self.fix_with_jumps(aaa.arcs):
+        for l1, l2 in arcs:
             fl1 = self.first_line(l1)
             fl2 = self.first_line(l2)
             if fl1 != fl2:
@@ -312,20 +316,41 @@ def _analyze_ast(self) -> None:
         self._missing_arc_fragments = aaa.missing_arc_fragments
 
     def fix_with_jumps(self, arcs: Iterable[TArc]) -> set[TArc]:
-        """Adjust arcs to fix jumps leaving `with` statements."""
+        """Adjust arcs to fix jumps leaving `with` statements.
+
+        Consider this code:
+
+            with open("/tmp/test", "w") as f1:
+                a = 2
+                b = 3
+            print(4)
+
+        In 3.10+, we get traces for lines 1, 2, 3, 1, 4.  But we want to present
+        it to the user as if it had been 1, 2, 3, 4.  The arc 3->1 should be
+        replaced with 3->4, and 1->4 should be removed.
+
+        For this code, the fixers dict is {(3, 1): ((1, 4), (3, 4))}.  The key
+        is the actual measured arc from the end of the with block back to the
+        start of the with-statement.  The values are start_next (the with
+        statement to the next statement after the with), and end_next (the end
+        of the with-statement to the next statement after the with).
+
+        With nested with-statements, we have to trace through a few levels to
+        correct a longer chain of arcs.
+
+        """
         to_remove = set()
         to_add = set()
         for arc in arcs:
             if arc in self._with_jump_fixers:
-                start = arc[0]
+                end0 = arc[0]
                 to_remove.add(arc)
-                start_next, prev_next = self._with_jump_fixers[arc]
+                start_next, end_next = self._with_jump_fixers[arc]
                 while start_next in self._with_jump_fixers:
                     to_remove.add(start_next)
-                    start_next, prev_next = self._with_jump_fixers[start_next]
-                    to_remove.add(prev_next)
-                to_add.add((start, prev_next[1]))
-                to_remove.add(arc)
+                    start_next, end_next = self._with_jump_fixers[start_next]
+                    to_remove.add(end_next)
+                to_add.add((end0, end_next[1]))
                 to_remove.add(start_next)
         arcs = (set(arcs) | to_add) - to_remove
         return arcs
@@ -700,15 +725,12 @@ def analyze(self) -> None:
     def with_jump_fixers(self) -> dict[TArc, tuple[TArc, TArc]]:
         """Get a dict with data for fixing jumps out of with statements.
 
-        Returns a dict.  The keys are arcs leaving a with statement by jumping
+        Returns a dict.  The keys are arcs leaving a with-statement by jumping
         back to its start.  The values are pairs: first, the arc from the start
         to the next statement, then the arc that exits the with without going
         to the start.
 
         """
-        if not env.PYBEHAVIOR.exit_through_with:
-            return {}
-
         fixers = {}
         with_nexts = {
             arc
@@ -721,9 +743,9 @@ def with_jump_fixers(self) -> dict[TArc, tuple[TArc, TArc]]:
                 continue
             assert len(nexts) == 1, f"Expected one arc, got {nexts} with {start = }"
             nxt = nexts.pop()
-            prvs = {arc[0] for arc in self.with_exits if arc[1] == start}
-            for prv in prvs:
-                fixers[(prv, start)] = ((start, nxt), (prv, nxt))
+            ends = {arc[0] for arc in self.with_exits if arc[1] == start}
+            for end in ends:
+                fixers[(end, start)] = ((start, nxt), (end, nxt))
         return fixers
 
     # Code object dispatchers: _code_object__*