Forward *args/**kwargs from functools.partial to the wrapper

PiperOrigin-RevId: 821720894

Author: superbobry

pytype/abstract/_function_base.py

@@ -192,6 +192,7 @@ def call(
           self.func.__self__  # pytype: disable=attribute-error
       )
     args = args.simplify(node, self.ctx, match_signature=sig)
+    del sig
     posargs = [u.AssignToNewVariable(node) for u in args.posargs]
     namedargs = {
         k: u.AssignToNewVariable(node) for k, u in args.namedargs.items()

pytype/overlays/functools_overlay.py

@@ -55,7 +55,12 @@ def new_slot(
         self.ctx,
         node,
         new.to_variable(node),
-        function.Args((cls, *args), kwargs),
+        function.Args(
+            (cls, *args),
+            kwargs,
+            call_context.starargs,
+            call_context.starstarargs,
+        ),
         fallback_to_unsolvable=False,
     )
     [specialized_obj] = specialized_obj.data
@@ -66,7 +71,7 @@ def new_slot(
     return node, obj.to_variable(node)
 
   def get_own_new(self, node, value) -> tuple[cfg.CFGNode, cfg.Variable]:
-    new = abstract.NativeFunction("__new__", self.new_slot, self.ctx)
+    new = NativeFunction("__new__", self.new_slot, self.ctx)
     return node, new.to_variable(node)
 
 
@@ -76,6 +81,8 @@ def bind_partial(node, cls, args, kwargs, ctx) -> BoundPartial:
   obj.underlying = args[0]
   obj.args = args[1:]
   obj.kwargs = kwargs
+  obj.starargs = call_context.starargs
+  obj.starstarargs = call_context.starstarargs
   return obj
 
 
@@ -115,6 +122,21 @@ def call(
   ) -> tuple[cfg.CFGNode, cfg.Variable]:
     # ``NativeFunction.call`` does not forward *args and **kwargs to the
     # underlying function, so we do it here to avoid changing core pytype APIs.
+    #
+    # The simplification below ensures that the *args/**kwargs cannot in fact
+    # be split into individual arguments. This logic follow the implementation
+    # in the base class.
+    sig = None
+    if isinstance(
+        self.func.__self__,  # pytype: disable=attribute-error
+        abstract.CallableClass,
+    ):
+      sig = function.Signature.from_callable(
+          self.func.__self__  # pytype: disable=attribute-error
+      )
+    args = args.simplify(node, self.ctx, match_signature=sig)
+    del sig
+
     starargs = args.starargs
     starstarargs = args.starstarargs
     if starargs is not None:
@@ -131,6 +153,8 @@ class BoundPartial(abstract.Instance, mixin.HasSlots):
   underlying: cfg.Variable
   args: tuple[cfg.Variable, ...]
   kwargs: dict[str, cfg.Variable]
+  starargs: cfg.Variable | None
+  starstarargs: cfg.Variable | None
 
   def __init__(self, ctx, cls, container=None):
     super().__init__(cls, ctx, container)
@@ -141,28 +165,72 @@ def __init__(self, ctx, cls, container=None):
 
   def get_signatures(self) -> Sequence[function.Signature]:
     sigs = []
-    args = function.Args(posargs=self.args, namedargs=self.kwargs)
+    args = function.Args(
+        self.args, self.kwargs, self.starargs, self.starstarargs
+    )
     for data in self.underlying.data:
       for sig in function.get_signatures(data):
         # Use the partial arguments as defaults in the signature, making them
         # optional but overwritable.
         defaults = sig.defaults.copy()
+        kwonly_params = [*sig.kwonly_params]
+        bound_param_names = set()
+        pos_only_count = sig.posonly_count
         for name, value, _ in sig.iter_args(args):
-          if value is not None:
-            defaults[name] = value
-        sigs.append(sig._replace(defaults=defaults))
+          if value is None:
+            continue
+          if sig.param_names.index(name) < sig.posonly_count:
+            # The parameter is positional-only, meaning that it cannot be
+            # overwritten via a keyword argument. Remove it.
+            bound_param_names.add(name)
+            sig.posonly_count -= 1
+            continue
+          if name not in sig.kwonly_params:
+            # The parameter can be overwritten via a keyword argument. Note
+            # that we still have to remove it from ``param_names`` to make
+            # sure it cannot be bound by position.
+            bound_param_names.add(name)
+            kwonly_params.append(name)
+
+          defaults[name] = value
+
+        sigs.append(
+            sig._replace(
+                param_names=tuple(
+                    n for n in sig.param_names if n not in bound_param_names
+                ),
+                posonly_count=pos_only_count,
+                kwonly_params=tuple(kwonly_params),
+                defaults=defaults,
+            )
+        )
     return sigs
 
   def call_slot(self, node: cfg.CFGNode, *args, **kwargs):
+    if self.starargs and call_context.starargs:
+      combined_starargs = self.ctx.convert.build_tuple(
+          node, (self.starargs, call_context.starargs)
+      )
+    else:
+      combined_starargs = call_context.starargs or self.starargs
+
+    if self.starstarargs and call_context.starstarargs:
+      d = abstract.Dict(self.ctx)
+      d.update(node, self.starstarargs.data[0])
+      d.update(node, call_context.starstarargs.data[0])
+      combined_starstarargs = d.to_variable(node)
+    else:
+      combined_starstarargs = call_context.starstarargs or self.starstarargs
+
     return function.call_function(
         self.ctx,
         node,
         self.underlying,
         function.Args(
             (*self.args, *args),
             {**self.kwargs, **kwargs},
-            call_context.starargs,
-            call_context.starstarargs,
+            combined_starargs,
+            combined_starstarargs,
         ),
         fallback_to_unsolvable=False,
     )

pytype/tests/test_attr2.py

@@ -230,6 +230,20 @@ class Foo:
       assert_type(foo.x, str)
     """)
 
+  def test_partial_with_positional_args_as_converter(self):
+    self.Check("""
+      import attr
+      import functools
+      def f(x: str, y: int) -> int:
+        del x
+        return y
+      @attr.s
+      class Foo:
+        x = attr.ib(converter=functools.partial(f, "foo"))
+      foo = Foo(x=0)
+      assert_type(foo.x, int)
+    """)
+
   def test_partial_overloaded_as_converter(self):
     self.Check("""
       import attr

pytype/tests/test_functions1.py

@@ -1079,6 +1079,39 @@ def f(a, b=None):
       partial_f(0)
     """)
 
+  def test_functools_partial_starstar(self):
+    self.Check("""
+      import functools
+      def f(*, a: str, b: int):
+        pass
+
+      def test(**kwargs):
+        partial_f = functools.partial(f, **kwargs)
+        partial_f()
+    """)
+
+  def test_functools_partial_called_with_starstar(self):
+    self.Check("""
+      import functools
+      def f(a: str, b: int, c: list):
+        pass
+      partial_f = functools.partial(f, "foo")
+
+      def test(**kwargs):
+        partial_f(42, **kwargs)
+    """)
+
+  def test_functools_starstar_everywhere(self):
+    self.Check("""
+      import functools
+      def f(*, a: str, b: int):
+        pass
+
+      def test(**kwargs):
+        partial_f = functools.partial(f, **kwargs)
+        partial_f(**kwargs)
+    """)
+
   def test_functools_partial_overloaded(self):
     self.Check("""
       import functools