Python: CG trace: Handle CALL_FUNCTION_EX

RasmusWL · RasmusWL · commit 0d05d96b504a · 2020-07-21T22:54:45.000+02:00
diff --git a/python/tools/recorded-call-graph-metrics/example/CALL_FUNCTION_EX.py b/python/tools/recorded-call-graph-metrics/example/CALL_FUNCTION_EX.py
@@ -0,0 +1,15 @@
+def func(*args, **kwargs):
+    print("func", args, kwargs)
+
+
+args = [1, 2, 3]
+kwargs = {"a": 1, "b": 2}
+
+# These gives rise to a CALL_FUNCTION_EX
+func(*args)
+func(**kwargs)
+func(*args, **kwargs)
+
+
+func(*args, foo="foo")
+func(*args, foo="foo", **kwargs)
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
@@ -162,13 +162,23 @@ def expr_from_instruction(instructions: List[Instruction], index: int) -> Byteco
         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"]:
+    elif inst.opname in [
+        "CALL_FUNCTION",
+        "CALL_METHOD",
+        "CALL_FUNCTION_KW",
+        "CALL_FUNCTION_EX",
+    ]:
         assert index > 0
         assert isinstance(inst.arg, int)
         if inst.opname in ["CALL_FUNCTION", "CALL_METHOD"]:
             num_stack_elems = inst.arg
         elif inst.opname == "CALL_FUNCTION_KW":
             num_stack_elems = inst.arg + 1
+        elif inst.opname == "CALL_FUNCTION_EX":
+            # top of stack _can_ be keyword argument dictionary (indicated by lowest bit
+            # set), always followed by the positional arguments (also if there are not
+            # any).
+            num_stack_elems = (1 if inst.arg & 1 == 1 else 0) + 1
 
         func_expr = expr_that_added_elem_to_stack(
             instructions, index - 1, num_stack_elems
