Properly handle break/continue/raise/return from except/else clauses

2 files changed, 121 insertions, 20 deletions

diff --git a/coverage/parser.py b/coverage/parser.py
index 307b83e6..756ec680 100644
--- a/coverage/parser.py
+++ b/coverage/parser.py
@@ -454,7 +454,7 @@ class AstArcAnalyzer(object):
             if isinstance(block, LoopBlock):
                 block.break_exits.update(exits)
                 break
-            elif isinstance(block, TryBlock) and block.final_start:
+            elif isinstance(block, TryBlock) and block.final_start is not None:
                 block.break_from.update(exits)
                 break
 
@@ -465,7 +465,7 @@ class AstArcAnalyzer(object):
                 for xit in exits:
                     self.arcs.add((xit, block.start))
                 break
-            elif isinstance(block, TryBlock) and block.final_start:
+            elif isinstance(block, TryBlock) and block.final_start is not None:
                 block.continue_from.update(exits)
                 break
 
@@ -473,11 +473,11 @@ class AstArcAnalyzer(object):
         """Add arcs due to jumps from `exits` being raises."""
         for block in self.nearest_blocks():
             if isinstance(block, TryBlock):
-                if block.handler_start:
+                if block.handler_start is not None:
                     for xit in exits:
                         self.arcs.add((xit, block.handler_start))
                     break
-                elif block.final_start:
+                elif block.final_start is not None:
                     block.raise_from.update(exits)
                     break
             elif isinstance(block, FunctionBlock):
@@ -488,7 +488,7 @@ class AstArcAnalyzer(object):
     def process_return_exits(self, exits):
         """Add arcs due to jumps from `exits` being returns."""
         for block in self.nearest_blocks():
-            if isinstance(block, TryBlock) and block.final_start:
+            if isinstance(block, TryBlock) and block.final_start is not None:
                 block.return_from.update(exits)
                 break
             elif isinstance(block, FunctionBlock):
@@ -568,9 +568,6 @@ class AstArcAnalyzer(object):
         return set()
 
     def _handle__Try(self, node):
-        # try/finally is tricky. If there's a finally clause, then we need a
-        # FinallyBlock to track what flows might go through the finally instead
-        # of their normal flow.
         if node.handlers:
             handler_start = self.line_for_node(node.handlers[0])
         else:
@@ -581,13 +578,27 @@ class AstArcAnalyzer(object):
         else:
             final_start = None
 
-        self.block_stack.append(TryBlock(handler_start=handler_start, final_start=final_start))
+        try_block = TryBlock(handler_start=handler_start, final_start=final_start)
+        self.block_stack.append(try_block)
 
         start = self.line_for_node(node)
         exits = self.add_body_arcs(node.body, from_line=start)
 
-        try_block = self.block_stack.pop()
+        # We're done with the `try` body, so this block no longer handles
+        # exceptions. We keep the block so the `finally` clause can pick up
+        # flows from the handlers and `else` clause.
+        if node.finalbody:
+            try_block.handler_start = None
+            if node.handlers:
+                # If there are `except` clauses, then raises in the try body
+                # will already jump to them.  Start this set over for raises in
+                # `except` and `else`.
+                try_block.raise_from = set([])
+        else:
+            self.block_stack.pop()
+
         handler_exits = set()
+
         last_handler_start = None
         if node.handlers:
             for handler_node in node.handlers:
@@ -608,20 +619,23 @@ class AstArcAnalyzer(object):
             exits = self.add_body_arcs(node.orelse, prev_lines=exits)
 
         exits |= handler_exits
+
         if node.finalbody:
+            self.block_stack.pop()
             final_from = (                  # You can get to the `finally` clause from:
                 exits |                         # the exits of the body or `else` clause,
-                try_block.break_from |          # or a `break` in the body,
-                try_block.continue_from |       # or a `continue` in the body,
-                try_block.return_from           # or a `return` in the body.
+                try_block.break_from |          # or a `break`,
+                try_block.continue_from |       # or a `continue`,
+                try_block.raise_from |          # or a `raise`,
+                try_block.return_from           # or a `return`.
             )
-            if node.handlers and last_handler_start is not None:
-                # If there was an "except X:" clause, then a "raise" in the
-                # body goes to the "except X:" before the "finally", but the
-                # "except" go to the finally.
-                final_from.add(last_handler_start)
-            else:
-                final_from |= try_block.raise_from
+            if node.handlers:
+                if last_handler_start is not None:
+                    # If we had handlers, and we didn't have a bare `except:`
+                    # handler, then the last handler jumps to the `finally` for the
+                    # unhandled exceptions.
+                    final_from.add(last_handler_start)
+
             exits = self.add_body_arcs(node.finalbody, prev_lines=final_from)
             if try_block.break_from:
                 self.process_break_exits(exits)
@@ -631,6 +645,7 @@ class AstArcAnalyzer(object):
                 self.process_raise_exits(exits)
             if try_block.return_from:
                 self.process_return_exits(exits)
+
         return exits
 
     def _handle__TryExcept(self, node):
diff --git a/tests/test_arcs.py b/tests/test_arcs.py
index 4513d085..5155264a 100644
--- a/tests/test_arcs.py
+++ b/tests/test_arcs.py
@@ -774,6 +774,92 @@ class ExceptionArcTest(CoverageTest):
             arcz=".1 12 28 89 9.  .3 34 46 6-2",
         )
 
+    def test_except_jump_finally(self):
+        self.check_coverage("""\
+            def func(x):
+                a = f = g = 2
+                try:
+                    for i in range(4):
+                        try:
+                            6/0
+                        except ZeroDivisionError:
+                            if x == 'break':
+                                a = 9
+                                break
+                            elif x == 'continue':
+                                a = 12
+                                continue
+                            elif x == 'return':
+                                a = 15                      # F
+                                return a, f, g, i           # G
+                            elif x == 'raise':              # H
+                                a = 18                      # I
+                                raise ValueError()          # J
+                        finally:
+                            f = 21                          # L
+                except ValueError:                          # M
+                    g = 23                                  # N
+                return a, f, g, i                           # O
+
+            assert func('break') == (9, 21, 2, 0)           # Q
+            assert func('continue') == (12, 21, 2, 3)       # R
+            assert func('return') == (15, 2, 2, 0)          # S
+            assert func('raise') == (18, 21, 23, 0)         # T
+            """,
+            arcz=
+                ".1 1Q QR RS ST T. "
+                ".2 23 34 45 56 4O 6L 7L "
+                "78 89 9A AL  8B BC CD DL  BE EF FG GL  EH HI IJ JL  HL "
+                "LO L4 L. LM "
+                "MN NO O.",
+            arcz_missing="6L 7L HL",
+            arcz_unpredicted="67",
+        )
+
+    def test_else_jump_finally(self):
+        self.check_coverage("""\
+            def func(x):
+                a = f = g = 2
+                try:
+                    for i in range(4):
+                        try:
+                            b = 6
+                        except ZeroDivisionError:
+                            pass
+                        else:
+                            if x == 'break':
+                                a = 11
+                                break
+                            elif x == 'continue':
+                                a = 14
+                                continue
+                            elif x == 'return':
+                                a = 17                      # H
+                                return a, f, g, i           # I
+                            elif x == 'raise':              # J
+                                a = 20                      # K
+                                raise ValueError()          # L
+                        finally:
+                            f = 23                          # N
+                except ValueError:                          # O
+                    g = 25                                  # P
+                return a, f, g, i                           # Q
+
+            assert func('break') == (11, 23, 2, 0)          # S
+            assert func('continue') == (14, 23, 2, 3)       # T
+            assert func('return') == (17, 2, 2, 0)          # U
+            assert func('raise') == (20, 23, 25, 0)         # V
+            """,
+            arcz=
+                ".1 1S ST TU UV V. "
+                ".2 23 34 45 56 6A 78 7N 8N 4Q "
+                "AB BC CN  AD DE EF FN  DG GH HI IN  GJ JK KL LN  JN "
+                "NQ N4 N. NO "
+                "OP PQ Q.",
+            arcz_missing="78 8N 7N JN",
+            arcz_unpredicted="",
+        )
+
 
 class YieldTest(CoverageTest):
     """Arc tests for generators."""