Fix python try except instrumentation

Author: jmthomas

openc3-cosmos-script-runner-api/scripts/script_instrumentor.py

@@ -73,10 +73,9 @@ def track_enter_leave(self, node):
         in_try = self.in_try
         if not in_try and type(node) in self.try_nodes:
             self.in_try = True
+
         # Visit the children of the node
         node = self.generic_visit(node)
-        if not in_try and type(node) in self.try_nodes:
-            self.in_try = False
         # ast.parse returns a module, so we need to extract
         # the first element of the body which is the node
         pre_line = ast.parse(
@@ -136,9 +135,10 @@ def track_enter_leave(self, node):
     # Call the pre_line_instrumentation ONLY and then execute the node
     def track_reached(self, node):
         # Determine if we're in a try block, this is used by track_enter_leave
-        in_try = self.in_try
-        if not in_try and type(node) in self.try_nodes:
+        if not self.in_try and type(node) in self.try_nodes:
             self.in_try = True
+        else:
+            self.in_try = False
 
         # Visit the children of the node
         node = self.generic_visit(node)

openc3-cosmos-script-runner-api/test/scripts/test_script_instrumentor.py

@@ -191,6 +191,36 @@ def test_exception_script(mock_running_script):
     ]
 
 
+def test_raise_with_try(mock_running_script):
+    script = """
+raise RuntimeError("Error")
+try:
+  pass
+except RuntimeError:
+  pass
+raise RuntimeError("Error")
+"""
+    parsed = ast.parse(script)
+    tree = ScriptInstrumentor("testfile.py").visit(parsed)
+    compiled = compile(tree, filename="testfile.py", mode="exec")
+    exec(compiled, {"RunningScript": mock_running_script})
+
+    assert mock_running_script.pre_lines == [
+        ("testfile.py", 2),
+        ("testfile.py", 3),
+        ("testfile.py", 4),
+        ("testfile.py", 7),
+    ]
+    assert mock_running_script.post_lines == [
+        ("testfile.py", 2),
+        ("testfile.py", 7),
+    ]
+    assert mock_running_script.exceptions == [
+        ("testfile.py", 2),
+        ("testfile.py", 7),
+    ]
+
+
 def test_import_future_script(mock_running_script):
     script = "from __future__ import annotations\nprint('hi')"
     parsed = ast.parse(script)