Python: CG trace: Handle subscript

RasmusWL · RasmusWL · commit 61b1d3eef3b0 · 2020-07-21T21:45:53.000+02:00
diff --git a/python/tools/recorded-call-graph-metrics/example/__getitem__.py b/python/tools/recorded-call-graph-metrics/example/__getitem__.py
@@ -0,0 +1,7 @@
+class Foo:
+    def __getitem__(self, key):
+        print("__getitem__")
+
+
+foo = Foo()
+foo["key"]  # this is recorded as a call :)
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
@@ -45,6 +45,15 @@ def __str__(self):
         return f"{self.object}.{self.attr_name}"
 
 
+@dataclasses.dataclass(frozen=True, eq=True, order=True)
+class BytecodeSubscript(BytecodeExpr):
+    key: BytecodeExpr
+    object: BytecodeExpr
+
+    def __str__(self):
+        return f"{self.object}[{self.key}]"
+
+
 @dataclasses.dataclass(frozen=True, eq=True, order=True)
 class BytecodeCall(BytecodeExpr):
     function: BytecodeExpr
@@ -147,6 +156,11 @@ def expr_from_instruction(instructions: List[Instruction], index: int) -> Byteco
         obj_expr = expr_that_added_elem_to_stack(instructions, index - 1, 0)
         return BytecodeAttribute(attr_name=attr_name, object=obj_expr)
 
+    elif inst.opname in ["BINARY_SUBSCR"]:
+        key_expr = expr_that_added_elem_to_stack(instructions, index - 1, 0)
+        obj_expr = expr_that_added_elem_to_stack(instructions, index - 1, 1)
+        return BytecodeSubscript(key=key_expr, object=obj_expr)
+
     # https://docs.python.org/3/library/dis.html#opcode-CALL_FUNCTION
     elif inst.opname in ["CALL_FUNCTION", "CALL_METHOD", "CALL_FUNCTION_KW"]:
         assert index > 0
