fix: in Python 3.10, leaving a with block exits through the with statement.

This need 3.10.0b3 (not yet released) to fully pass.

4 files changed, 185 insertions, 16 deletions

diff --git a/coverage/env.py b/coverage/env.py
index cc8ca8b7..81f61794 100644
--- a/coverage/env.py
+++ b/coverage/env.py
@@ -99,6 +99,9 @@ class PYBEHAVIOR:
     # Are "if 0:" lines (and similar) kept in the compiled code?
     keep_constant_test = pep626
 
+    # When leaving a with-block, do we visit the with-line again for the exit?
+    exit_through_with = (PYVERSION >= (3, 10, 0, 'beta'))
+
 # Coverage.py specifics.
 
 # Are we using the C-implemented trace function?
diff --git a/coverage/parser.py b/coverage/parser.py
index 1c8ecc3e..ff395dad 100644
--- a/coverage/parser.py
+++ b/coverage/parser.py
@@ -513,8 +513,8 @@ class TryBlock(BlockBase):
         # that need to route through the "finally:" clause.
         self.break_from = set()
         self.continue_from = set()
-        self.return_from = set()
         self.raise_from = set()
+        self.return_from = set()
 
     def process_break_exits(self, exits, add_arc):
         if self.final_start is not None:
@@ -532,11 +532,10 @@ class TryBlock(BlockBase):
         if self.handler_start is not None:
             for xit in exits:
                 add_arc(xit.lineno, self.handler_start, xit.cause)
-            return True
-        elif self.final_start is not None:
+        else:
+            assert self.final_start is not None
             self.raise_from.update(exits)
-            return True
-        return False
+        return True
 
     def process_return_exits(self, exits, add_arc):
         if self.final_start is not None:
@@ -545,6 +544,44 @@ class TryBlock(BlockBase):
         return False
 
 
+class WithBlock(BlockBase):
+    """A block on the block stack representing a `with` block."""
+    @contract(start=int)
+    def __init__(self, start):
+        # We only ever use this block if it is needed, so that we don't have to
+        # check this setting in all the methods.
+        assert env.PYBEHAVIOR.exit_through_with
+
+        # The line number of the with statement.
+        self.start = start
+
+        # The ArcStarts for breaks/continues/returns/raises inside the "with:"
+        # that need to go through the with-statement while exiting.
+        self.break_from = set()
+        self.continue_from = set()
+        self.raise_from = set()
+        self.return_from = set()
+
+    def _process_exits(self, exits, add_arc, from_set):
+        """Helper to process the four kinds of exits."""
+        for xit in exits:
+            add_arc(xit.lineno, self.start, xit.cause)
+        from_set.update(exits)
+        return True
+
+    def process_break_exits(self, exits, add_arc):
+        return self._process_exits(exits, add_arc, self.break_from)
+
+    def process_continue_exits(self, exits, add_arc):
+        return self._process_exits(exits, add_arc, self.continue_from)
+
+    def process_raise_exits(self, exits, add_arc):
+        return self._process_exits(exits, add_arc, self.raise_from)
+
+    def process_return_exits(self, exits, add_arc):
+        return self._process_exits(exits, add_arc, self.return_from)
+
+
 class ArcStart(collections.namedtuple("Arc", "lineno, cause")):
     """The information needed to start an arc.
 
@@ -731,7 +768,7 @@ class AstArcAnalyzer:
             # statement), or it's something we overlooked.
             if env.TESTING:
                 if node_name not in self.OK_TO_DEFAULT:
-                    raise Exception(f"*** Unhandled: {node}")
+                    raise Exception(f"*** Unhandled: {node}")       # pragma: only failure
 
             # Default for simple statements: one exit from this node.
             return {ArcStart(self.line_for_node(node))}
@@ -865,14 +902,14 @@ class AstArcAnalyzer:
     @contract(exits='ArcStarts')
     def process_break_exits(self, exits):
         """Add arcs due to jumps from `exits` being breaks."""
-        for block in self.nearest_blocks():
+        for block in self.nearest_blocks():                         # pragma: always breaks
             if block.process_break_exits(exits, self.add_arc):
                 break
 
     @contract(exits='ArcStarts')
     def process_continue_exits(self, exits):
         """Add arcs due to jumps from `exits` being continues."""
-        for block in self.nearest_blocks():
+        for block in self.nearest_blocks():                         # pragma: always breaks
             if block.process_continue_exits(exits, self.add_arc):
                 break
 
@@ -886,7 +923,7 @@ class AstArcAnalyzer:
     @contract(exits='ArcStarts')
     def process_return_exits(self, exits):
         """Add arcs due to jumps from `exits` being returns."""
-        for block in self.nearest_blocks():
+        for block in self.nearest_blocks():                         # pragma: always breaks
             if block.process_return_exits(exits, self.add_arc):
                 break
 
@@ -1014,6 +1051,9 @@ class AstArcAnalyzer:
         else:
             final_start = None
 
+        # This is true by virtue of Python syntax: have to have either except
+        # or finally, or both.
+        assert handler_start is not None or final_start is not None
         try_block = TryBlock(handler_start, final_start)
         self.block_stack.append(try_block)
 
@@ -1159,7 +1199,32 @@ class AstArcAnalyzer:
     @contract(returns='ArcStarts')
     def _handle__With(self, node):
         start = self.line_for_node(node)
+        if env.PYBEHAVIOR.exit_through_with:
+            self.block_stack.append(WithBlock(start=start))
         exits = self.add_body_arcs(node.body, from_start=ArcStart(start))
+        if env.PYBEHAVIOR.exit_through_with:
+            with_block = self.block_stack.pop()
+            with_exit = {ArcStart(start)}
+            if exits:
+                for xit in exits:
+                    self.add_arc(xit.lineno, start)
+                exits = with_exit
+            if with_block.break_from:
+                self.process_break_exits(
+                    self._combine_finally_starts(with_block.break_from, with_exit)
+                    )
+            if with_block.continue_from:
+                self.process_continue_exits(
+                    self._combine_finally_starts(with_block.continue_from, with_exit)
+                    )
+            if with_block.raise_from:
+                self.process_raise_exits(
+                    self._combine_finally_starts(with_block.raise_from, with_exit)
+                    )
+            if with_block.return_from:
+                self.process_return_exits(
+                    self._combine_finally_starts(with_block.return_from, with_exit)
+                    )
         return exits
 
     _handle__AsyncWith = _handle__With
diff --git a/metacov.ini b/metacov.ini
index bed5f940..9dab77aa 100644
--- a/metacov.ini
+++ b/metacov.ini
@@ -74,6 +74,8 @@ exclude_lines =
 
 partial_branches =
     pragma: part covered
+    # A for-loop that always hits its break statement
+    pragma: always breaks
     pragma: if failure
     pragma: part started
     if env.TESTING:
diff --git a/tests/test_arcs.py b/tests/test_arcs.py
index 905430e6..495a10f3 100644
--- a/tests/test_arcs.py
+++ b/tests/test_arcs.py
@@ -166,19 +166,43 @@ class WithTest(CoverageTest):
     """Arc-measuring tests involving context managers."""
 
     def test_with(self):
+        if env.PYBEHAVIOR.exit_through_with:
+            arcz = ".1 .2 23 34 42 2. 16 6."
+        else:
+            arcz = ".1 .2 23 34 4. 16 6."
         self.check_coverage("""\
             def example():
-                with open("test", "w") as f: # exit
-                    f.write("")
-                    return 1
+                with open("test", "w") as f:
+                    f.write("3")
+                    a = 4
+
+            example()
+            """,
+            arcz=arcz,
+            )
+
+    def test_with_return(self):
+        if env.PYBEHAVIOR.exit_through_with:
+            arcz = ".1 .2 23 34 42 2. 16 6."
+        else:
+            arcz = ".1 .2 23 34 4. 16 6."
+        self.check_coverage("""\
+            def example():
+                with open("test", "w") as f:
+                    f.write("3")
+                    return 4
 
             example()
             """,
-            arcz=".1 .2 23 34 4. 16 6."
+            arcz=arcz,
             )
 
     def test_bug_146(self):
         # https://github.com/nedbat/coveragepy/issues/146
+        if env.PYBEHAVIOR.exit_through_with:
+            arcz = ".1 12 23 32 24 41 15 5."
+        else:
+            arcz = ".1 12 23 34 41 15 5."
         self.check_coverage("""\
             for i in range(2):
                 with open("test", "w") as f:
@@ -186,7 +210,56 @@ class WithTest(CoverageTest):
                 print(4)
             print(5)
             """,
-            arcz=".1 12 23 34 41 15 5."
+            arcz=arcz,
+            )
+
+    def test_nested_with_return(self):
+        if env.PYBEHAVIOR.exit_through_with:
+            arcz = ".1 .2 23 34 45 56 64 42 2. 18 8."
+        else:
+            arcz = ".1 .2 23 34 45 56 6. 18 8."
+        self.check_coverage("""\
+            def example(x):
+                with open("test", "w") as f2:
+                    a = 3
+                    with open("test2", "w") as f4:
+                        f2.write("5")
+                        return 6
+
+            example(8)
+            """,
+            arcz=arcz,
+            )
+
+    def test_break_through_with(self):
+        if env.PYBEHAVIOR.exit_through_with:
+            arcz = ".1 12 23 34 42 25 15 5."
+        else:
+            arcz = ".1 12 23 34 45 15 5."
+        self.check_coverage("""\
+            for i in range(1+1):
+                with open("test", "w") as f:
+                    print(3)
+                    break
+            print(5)
+            """,
+            arcz=arcz,
+            arcz_missing="15",
+            )
+
+    def test_continue_through_with(self):
+        if env.PYBEHAVIOR.exit_through_with:
+            arcz = ".1 12 23 34 42 21 15 5."
+        else:
+            arcz = ".1 12 23 34 41 15 5."
+        self.check_coverage("""\
+            for i in range(1+1):
+                with open("test", "w") as f:
+                    print(3)
+                    continue
+            print(5)
+            """,
+            arcz=arcz,
             )
 
 
@@ -678,6 +751,26 @@ class ExceptionArcTest(CoverageTest):
             arcz_missing="3D BC CD",
         )
 
+    def test_break_continue_without_finally(self):
+        self.check_coverage("""\
+            a, c, d, i = 1, 1, 1, 99
+            try:
+                for i in range(3):
+                    try:
+                        a = 5
+                        if i > 0:
+                            break
+                        continue
+                    except:
+                        c = 10
+            except:
+                d = 12                              # C
+            assert a == 5 and c == 1 and d == 1     # D
+            """,
+            arcz=".1 12 23 34 3D 45 56 67 68 7D 83 9A A3 BC CD D.",
+            arcz_missing="3D 9A A3 BC CD",
+        )
+
     def test_continue_through_finally(self):
         if env.PYBEHAVIOR.finally_jumps_back:
             arcz = ".1 12 23 34 3D 45 56 67 68 73 7A 8A A3 A7 BC CD D."
@@ -1632,13 +1725,19 @@ class AsyncTest(CoverageTest):
         assert self.stdout() == "a\nb\nc\n.\n"
 
     def test_async_with(self):
+        if env.PYBEHAVIOR.exit_through_with:
+            arcz = ".1 1. .2 23 32 2."
+            arcz_missing = ".2 23 32 2."
+        else:
+            arcz = ".1 1. .2 23 3."
+            arcz_missing = ".2 23 3."
         self.check_coverage("""\
             async def go():
                 async with x:
                     pass
             """,
-            arcz=".1 1. .2 23 3.",
-            arcz_missing=".2 23 3.",
+            arcz=arcz,
+            arcz_missing=arcz_missing,
         )
 
     def test_async_decorator(self):