rewrite: Basic CALL_FUNCTION_EX support.

martindemello · rchen152 · commit 70f760992f89 · 2024-04-10T14:40:27.000-07:00
Supports unpacking a concrete iterable in a function call.

PiperOrigin-RevId: 623580582
diff --git a/pytype/rewrite/abstract/abstract.py b/pytype/rewrite/abstract/abstract.py
@@ -37,6 +37,7 @@
 Tuple = _containers.Tuple
 
 ConstKeyDict = _internal.ConstKeyDict
+FunctionArgTuple = _internal.FunctionArgTuple
 Splat = _internal.Splat
 
 get_atomic_constant = _utils.get_atomic_constant
diff --git a/pytype/rewrite/abstract/internal.py b/pytype/rewrite/abstract/internal.py
@@ -1,6 +1,6 @@
 """Abstract types used internally by pytype."""
 
-from typing import Any, Dict, Sequence
+from typing import Dict, Sequence, Tuple
 
 import immutabledict
 
@@ -30,6 +30,22 @@ def _attrs(self):
     return (immutabledict.immutabledict(self.constant),)
 
 
+class FunctionArgTuple(base.BaseValue):
+  """Representation of a function arg tuple."""
+
+  def __init__(self, ctx: base.ContextType, constant: Tuple[_Variable, ...]):
+    super().__init__(ctx)
+    assert isinstance(constant, tuple), constant
+    self.constant = constant
+
+  def __repr__(self):
+    return f"FunctionArgTuple({self.constant!r})"
+
+  @property
+  def _attrs(self):
+    return (self.constant,)
+
+
 class Splat(base.BaseValue):
   """Representation of unpacked iterables.
 
diff --git a/pytype/rewrite/frame.py b/pytype/rewrite/frame.py
@@ -571,6 +571,40 @@ def byte_CALL_FUNCTION(self, opcode):
     callargs = abstract.Args(posargs=tuple(args), frame=self)
     self._call_function(func, callargs)
 
+  def _unpack_starargs(self, starargs):
+    # TODO(b/331853896): This follows vm_utils.ensure_unpacked_starargs, but
+    # does not yet handle indefinite iterables.
+    posargs = starargs.get_atomic_value()
+    if isinstance(posargs, abstract.FunctionArgTuple):
+      # This has already been converted
+      pass
+    elif isinstance(posargs, abstract.Tuple):
+      posargs = abstract.FunctionArgTuple(self._ctx, posargs.constant)
+    elif isinstance(posargs, tuple):
+      posargs = abstract.FunctionArgTuple(self._ctx, posargs)
+    else:
+      assert False, f'unexpected posargs type: {posargs}: {type(posargs)}'
+    return posargs
+
+  def _unpack_starstarargs(self, starstarargs):
+    kwargs = abstract.get_atomic_constant(starstarargs, dict)
+    return {abstract.get_atomic_constant(k, str): v
+            for k, v in kwargs.items()}
+
+  def byte_CALL_FUNCTION_EX(self, opcode):
+    if opcode.arg & _Flags.CALL_FUNCTION_EX_HAS_KWARGS:
+      starstarargs = self._stack.pop()
+      kwargs = self._unpack_starstarargs(starstarargs)
+    else:
+      kwargs = _EMPTY_MAP
+    starargs = self._stack.pop()
+    posargs = self._unpack_starargs(starargs).constant
+    func = self._stack.pop()
+    if self._code.python_version >= (3, 11):
+      self._stack.pop_and_discard()
+    callargs = abstract.Args(posargs=posargs, kwargs=kwargs, frame=self)
+    self._call_function(func, callargs)
+
   def byte_CALL_METHOD(self, opcode):
     args = self._stack.popn(opcode.arg)
     func = self._stack.pop()
diff --git a/pytype/rewrite/frame_test.py b/pytype/rewrite/frame_test.py
@@ -635,6 +635,16 @@ def f(x, *, y):
     self.assertConstantVar(callargs.posargs[0], 1)
     self.assertConstantVar(callargs.kwargs['y'], 2)
 
+  @test_utils.skipBeforePy((3, 11), 'Relies on 3.11+ bytecode')
+  def test_call_function_ex_no_crash(self):
+    frame = self._make_frame("""
+      def f(x, y, z):
+        pass
+      a = (1, 2)
+      f(*a, z=3)
+    """)
+    frame.run()
+
 
 if __name__ == '__main__':
   unittest.main()
