Python: CG trace: Handle BUILD_TUPLE

RasmusWL · RasmusWL · commit 8c8656ccca2c · 2020-07-21T23:05:49.000+02:00
diff --git a/python/tools/recorded-call-graph-metrics/example/BUILD_TUPLE.py b/python/tools/recorded-call-graph-metrics/example/BUILD_TUPLE.py
@@ -0,0 +1,9 @@
+def foo():
+    print("foo")
+
+
+def bar():
+    print("bar")
+
+
+(foo, bar)[0]()
diff --git a/python/tools/recorded-call-graph-metrics/src/cg_trace/bytecode_reconstructor.py b/python/tools/recorded-call-graph-metrics/src/cg_trace/bytecode_reconstructor.py
@@ -54,6 +54,19 @@ def __str__(self):
         return f"{self.object}[{self.key}]"
 
 
+@dataclasses.dataclass(frozen=True, eq=True, order=True)
+class BytecodeTuple(BytecodeExpr):
+    elements: List[BytecodeExpr]
+
+    def __str__(self):
+        elements_formatted = (
+            ", ".join(str(e) for e in self.elements)
+            if len(self.elements) > 1
+            else f"{self.elements[0]},"
+        )
+        return f"({elements_formatted})"
+
+
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeCall(BytecodeExpr):
     function: BytecodeExpr
@@ -161,6 +174,14 @@ def expr_from_instruction(instructions: List[Instruction], index: int) -> Byteco
         obj_expr = expr_that_added_elem_to_stack(instructions, index - 1, 1)
         return BytecodeSubscript(key=key_expr, object=obj_expr)
 
+    elif inst.opname in ["BUILD_TUPLE"]:
+        elements = []
+        for i in range(inst.arg):
+            element_expr = expr_that_added_elem_to_stack(instructions, index - 1, i)
+            elements.append(element_expr)
+        elements.reverse()
+        return BytecodeTuple(elements=elements)
+
     # https://docs.python.org/3/library/dis.html#opcode-CALL_FUNCTION
     elif inst.opname in [
         "CALL_FUNCTION",
