Fixed issue with elif

src/dynapyt/instrument/CodeInstrumenter.py

@@ -1606,19 +1606,22 @@ def leave_If(self, original_node, updated_node):
                     + [cst.SimpleStatementLine(body=[cst.Expr(value=end_call)])]
                 )
 
-            if updated_node.orelse is None:
-                new_orelse_body = [
-                    cst.SimpleStatementLine(body=[cst.Expr(value=end_call)])
-                ]
-            elif m.matches(updated_node.orelse.body, m.SimpleStatementSuite()):
-                new_orelse_body = list(updated_node.orelse.body.body) + [
-                    cst.SimpleStatementLine(body=[cst.Expr(value=end_call)])
-                ]
+            if not m.matches(updated_node.orelse, m.If()):
+                if updated_node.orelse is None:
+                    new_orelse_body = [
+                        cst.SimpleStatementLine(body=[cst.Expr(value=end_call)])
+                    ]
+                elif m.matches(updated_node.orelse.body, m.SimpleStatementSuite()):
+                    new_orelse_body = list(updated_node.orelse.body.body) + [
+                        cst.SimpleStatementLine(body=[cst.Expr(value=end_call)])
+                    ]
+                else:
+                    new_orelse_body = list(updated_node.orelse.body.body) + [
+                        cst.SimpleStatementLine(body=[cst.Expr(value=end_call)])
+                    ]
+                new_orelse = cst.Else(body=cst.IndentedBlock(body=new_orelse_body))
             else:
-                new_orelse_body = list(updated_node.orelse.body.body) + [
-                    cst.SimpleStatementLine(body=[cst.Expr(value=end_call)])
-                ]
-            new_orelse = cst.Else(body=cst.IndentedBlock(body=new_orelse_body))
+                new_orelse = updated_node.orelse
         else:
             new_body = updated_node.body
             new_orelse = updated_node.orelse

tests/trace_single_hook/elif/analysis.py

@@ -0,0 +1,23 @@
+from typing import Any, Optional
+from dynapyt.analyses.BaseAnalysis import BaseAnalysis
+from dynapyt.instrument.IIDs import IIDs
+
+
+class TestAnalysis(BaseAnalysis):
+    def begin_execution(self) -> None:
+        print("begin execution")
+
+    def enter_control_flow(self, dyn_ast: str, iid: int, cond_value: bool) -> Optional[bool]:
+        print(f"enter control flow event with condition {cond_value}")
+
+    def enter_if(self, dyn_ast: str, iid: int, cond_value: bool) -> Optional[bool]:
+        print(f"if condition evaluates to {cond_value}")
+
+    def exit_if(self, dyn_ast: str, iid: int):
+        print(f"done with if statement")
+
+    def exit_control_flow(self, dyn_ast: str, iid: int) -> None:
+        print(f"done with control flow statement")
+
+    def end_execution(self) -> None:
+        print("end execution")

tests/trace_single_hook/elif/expected.txt

@@ -0,0 +1,11 @@
+begin execution
+enter control flow event with condition False
+if condition evaluates to False
+enter control flow event with condition False
+if condition evaluates to False
+enter control flow event with condition True
+if condition evaluates to True
+foo is equal to 50
+done with control flow statement
+done with if statement
+end execution

tests/trace_single_hook/elif/program.py

@@ -0,0 +1,9 @@
+foo = 45 + 5
+if foo > 50:
+    print("foo is greater than 50")
+elif foo == 45:
+    print("foo is equal to 45")
+elif foo == 50:
+    print("foo is equal to 50")
+else:
+    print("foo is less than 50")