- fixed 'next' computation for exception handlers;

plasma-umass · Nov 13, 2023 · 8e41a67 · 8e41a67
1 parent 1f42251
commit 8e41a67
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/src/slipcover/branch.py b/src/slipcover/branch.py
@@ -113,6 +113,10 @@ def visit_Match(self, node: ast.Match) -> ast.Match:
                 # each case continues after the 'match'
                 for item in field:
                     item.next_node = node.next_node
+            elif isinstance(node, try_type) and name == 'handlers':
+                # each 'except' continues either in 'finally', or after the 'try'
+                for h in field:
+                    h.next_node = node.finalbody[0] if node.finalbody else node.next_node
             elif isinstance(field, list):
                 # if a field is a list, each item but the last one continues with the next item
                 prev = None

diff --git a/tests/test_branch.py b/tests/test_branch.py
@@ -624,14 +624,19 @@ def foo(x):
                 if y < 0:
                     y = 0
             except{star} RuntimeException:
-                y = 0
+                if y < 2:
+                    y = 0
+            except{star} FileNotFoundError:
+                if y < 2:
+                    y = 0
+
             return 2*y
     """)
 
 
     t = br.preinstrument(t)
     code = compile(t, "foo", "exec")
-    assert [(4,5), (4,8)] == get_branches(code)
+    assert [(4,5), (4,13), (7,8), (7,13), (10,11), (10,13)] == get_branches(code)
 
 
 def test_try_finally():
@@ -660,15 +665,16 @@ def foo(x):
                 if y < 0:
                     y = 0
             except{star} RuntimeException:
-                y = -1
+                if y < 2:
+                    y = -1
             else:
                 y = 2*y
     """)
 
 
     t = br.preinstrument(t)
     code = compile(t, "foo", "exec")
-    assert [(4,5), (4,9)] == get_branches(code)
+    assert [(4,5), (4,10), (7,0), (7,8)] == get_branches(code)
 
 
 @pytest.mark.parametrize("star", ['', '*'] if PYTHON_VERSION >= (3,11) else [''])
@@ -680,7 +686,8 @@ def foo(x):
                 if y < 0:
                     y = 0
             except{star} RuntimeException:
-                y = -1
+                if y < 2:
+                    y = -1
             else:
                 if y > 5:
                     y = 42
@@ -691,4 +698,4 @@ def foo(x):
 
     t = br.preinstrument(t)
     code = compile(t, "foo", "exec")
-    assert [(4,5), (4,9), (9,10), (9,12)] == get_branches(code)
+    assert [(4,5), (4,10), (7,8), (7,13), (10,11), (10,13)] == get_branches(code)