Added an overlay for functools.partial

The goal of the overlay is to delay return type inference until the partial
object is called, allowing pytype to use all available arguments instead of
just the ones provided to `functools.partial`.

PiperOrigin-RevId: 748656406

Author: superbobry

pytype/overlays/CMakeLists.txt

@@ -161,6 +161,8 @@ py_library(
   DEPS
     .overlay
     .special_builtins
+    pytype.abstract.abstract
+    pytype.typegraph.cfg
 )
 
 py_library(

pytype/overlays/functools_overlay.py

@@ -1,7 +1,20 @@
 """Overlay for functools."""
 
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING
+
+from pytype.abstract import abstract
+from pytype.abstract import function
+from pytype.abstract import mixin
 from pytype.overlays import overlay
 from pytype.overlays import special_builtins
+from pytype.typegraph import cfg
+
+if TYPE_CHECKING:
+  from pytype import context  # pylint: disable=g-import-not-at-top
+
 
 _MODULE_NAME = "functools"
 
@@ -14,6 +27,80 @@ def __init__(self, ctx):
         "cached_property": overlay.add_name(
             "cached_property", special_builtins.Property.make_alias
         ),
+        "partial": Partial,
     }
     ast = ctx.loader.import_name(_MODULE_NAME)
     super().__init__(ctx, _MODULE_NAME, member_map, ast)
+
+
+class Partial(abstract.PyTDClass, mixin.HasSlots):
+  """Implementation of functools.partial."""
+
+  def __init__(self, ctx: "context.Context", module: str):
+    pytd_cls = ctx.loader.lookup_pytd(module, "partial")
+    super().__init__("partial", pytd_cls, ctx)
+    mixin.HasSlots.init_mixin(self)
+
+    self._pytd_new = self.pytd_cls.Lookup("__new__")
+
+  def new_slot(
+      self, node, cls, *args, **kwargs
+  ) -> tuple[cfg.CFGNode, cfg.Variable]:
+    # Make sure the call is well typed before binding the partial
+    new = self.ctx.convert.convert_pytd_function(self._pytd_new)
+    _, specialized_obj = function.call_function(
+        self.ctx,
+        node,
+        new.to_variable(node),
+        function.Args(posargs=(cls, *args), namedargs=kwargs),
+        fallback_to_unsolvable=False,
+    )
+    [specialized_obj] = specialized_obj.data
+    type_arg = specialized_obj.get_formal_type_parameter("_T")
+    [cls] = cls.data
+    cls = abstract.ParameterizedClass(cls, {"_T": type_arg}, self.ctx)
+    obj = bind_partial(node, cls, args, kwargs, self.ctx)
+    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)
+    return node, new.to_variable(node)
+
+
+def bind_partial(node, cls, args, kwargs, ctx) -> BoundPartial:
+  del node  # Unused.
+  obj = BoundPartial(ctx, cls)
+  obj.underlying = args[0]
+  obj.args = args[1:]
+  obj.kwargs = kwargs
+  return obj
+
+
+class BoundPartial(abstract.Instance, mixin.HasSlots):
+  """An instance of functools.partial."""
+
+  underlying: cfg.Variable
+  args: Sequence[cfg.Variable]
+  kwargs: Mapping[str, cfg.Variable]
+
+  def __init__(self, ctx, cls, container=None):
+    super().__init__(cls, ctx, container)
+    mixin.HasSlots.init_mixin(self)
+    self.set_native_slot("__call__", self.call_slot)
+
+  @property
+  def func(self) -> cfg.Variable:
+    # The ``func`` attribute marks this class as a wrapper for
+    # ``maybe_unwrap_decorated_function``.
+    return self.underlying
+
+  def call_slot(self, node: cfg.CFGNode, *args, **kwargs):
+    return function.call_function(
+        self.ctx,
+        node,
+        self.underlying,
+        function.Args(
+            posargs=(*self.args, *args), namedargs={**self.kwargs, **kwargs}
+        ),
+        fallback_to_unsolvable=False,
+    )

pytype/tests/test_functions1.py

@@ -1079,6 +1079,37 @@ def f(a, b=None):
       partial_f(0)
     """)
 
+  def test_functools_partial_overloaded1(self):
+    self.Check("""
+      import functools
+      from typing import overload
+      @overload
+      def f(x: int) -> int: ...
+      @overload
+      def f(x: str) -> str: ...
+      def f(x):
+        return x
+      partial_f = functools.partial(f)
+      # This should be functools.partial[int | str], but pytype does not
+      # seem to infer that.
+      assert_type(partial_f, functools.partial)
+      assert_type(partial_f(1), int)
+      assert_type(partial_f("s"), str)
+    """)
+
+  def test_functools_partial_overloaded2(self):
+    self.Check("""
+      import functools
+      from typing import overload
+      @overload
+      def f(x: int, y: int) -> int: ...
+      @overload
+      def f(x: str, y: str) -> str: ...
+      def f(x, y):
+        return x
+      partial_f = functools.partial(f, 42)
+    """)
+
   def test_functools_partial_class(self):
     self.Check("""
       import functools