fix: multi-line with-statements exit correctly. #1880

Author: nedbat

coverage/env.py

@@ -81,6 +81,10 @@ class PYBEHAVIOR:
     # When leaving a with-block, do we visit the with-line again for the exit?
     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?
+    exit_with_through_ctxmgr = (PYVERSION >= (3, 12))
+
     # Match-case construct.
     match_case = (PYVERSION >= (3, 10))
 

coverage/parser.py

@@ -1261,12 +1261,19 @@ def _handle__While(self, node: ast.While) -> set[ArcStart]:
         return exits
 
     def _handle__With(self, node: ast.With) -> set[ArcStart]:
-        start = self.line_for_node(node)
+        if env.PYBEHAVIOR.exit_with_through_ctxmgr:
+            starts = [self.line_for_node(item.context_expr) for item in node.items]
+        else:
+            starts = [self.line_for_node(node)]
         if env.PYBEHAVIOR.exit_through_with:
-            self.current_with_starts.add(start)
-            self.all_with_starts.add(start)
-        exits = self.process_body(node.body, from_start=ArcStart(start))
+            for start in starts:
+                self.current_with_starts.add(start)
+                self.all_with_starts.add(start)
+
+        exits = self.process_body(node.body, from_start=ArcStart(starts[-1]))
+
         if env.PYBEHAVIOR.exit_through_with:
+            start = starts[-1]
             self.current_with_starts.remove(start)
             with_exit = {ArcStart(start)}
             if exits:

tests/test_arcs.py

@@ -347,6 +347,41 @@ def test_with_with_lambda(self) -> None:
             branchz_missing="",
         )
 
+    def test_multiline_with(self) -> None:
+        # https://github.com/nedbat/coveragepy/issues/1880
+        self.check_coverage("""\
+            import contextlib, itertools
+            nums = itertools.count()
+            with (
+                contextlib.nullcontext() as x,
+            ):
+                while next(nums) < 6:
+                    y = 7
+            z = 8
+            """,
+            branchz="67 68",
+            branchz_missing="",
+        )
+
+
+    def test_multi_multiline_with(self) -> None:
+        # https://github.com/nedbat/coveragepy/issues/1880
+        self.check_coverage("""\
+            import contextlib, itertools
+            nums = itertools.count()
+            with (
+                contextlib.nullcontext() as x,
+                contextlib.nullcontext() as y,
+                contextlib.nullcontext() as z,
+            ):
+                while next(nums) < 8:
+                    y = 9
+            z = 10
+            """,
+            branchz="89 8A",
+            branchz_missing="",
+        )
+
 
 class LoopArcTest(CoverageTest):
     """Arc-measuring tests involving loops."""