Allowed functools.partial to be used as a converter= in attrs.field aka attr.ib

pytype can now do proper checking of such converter declarations.

PiperOrigin-RevId: 820705345

Author: superbobry

pytype/abstract/function.py

@@ -72,6 +72,22 @@ def get_signatures(func: "_function_base.Function") -> "list[Signature]":
     return []
   elif isinstance(func.cls, _abstract.CallableClass):
     return [Signature.from_callable(func.cls)]
+  elif isinstance(func, (_abstract.InterpreterClass, _abstract.PyTDClass)):
+    if isinstance(func, _abstract.PyTDClass) and "__init__" in func:
+      func.load_lazy_attribute("__init__")
+    if (init_var := func.members.get("__init__")) and len(init_var.data) == 1:
+      sigs = []
+      for sig in get_signatures(init_var.data[0]):
+        sig = sig.drop_first_parameter()  # drop "self"
+        sigs.append(
+            sig._replace(annotations=sig.annotations | {"return": func})
+        )
+      return sigs
+    # The class does not have __init__? Bail out!
+    # TODO(slebedev): Consider handling __new__ and metaclass.__call__ here.
+    return [Signature.from_any()]
+  elif hasattr(func, "get_signatures"):
+    return func.get_signatures()
   else:
     unwrapped = abstract_utils.maybe_unwrap_decorated_function(func)
     if unwrapped:

pytype/overlays/functools_overlay.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
 import threading
 from typing import Any, TYPE_CHECKING
 
@@ -129,8 +129,8 @@ class BoundPartial(abstract.Instance, mixin.HasSlots):
   """An instance of functools.partial."""
 
   underlying: cfg.Variable
-  args: Sequence[cfg.Variable]
-  kwargs: Mapping[str, cfg.Variable]
+  args: tuple[cfg.Variable, ...]
+  kwargs: dict[str, cfg.Variable]
 
   def __init__(self, ctx, cls, container=None):
     super().__init__(cls, ctx, container)
@@ -139,11 +139,19 @@ def __init__(self, ctx, cls, container=None):
         "__call__", NativeFunction("__call__", self.call_slot, self.ctx)
     )
 
-  @property
-  def func(self) -> cfg.Variable:
-    # The ``func`` attribute marks this class as a wrapper for
-    # ``maybe_unwrap_decorated_function``.
-    return self.underlying
+  def get_signatures(self) -> Sequence[function.Signature]:
+    sigs = []
+    args = function.Args(posargs=self.args, namedargs=self.kwargs)
+    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()
+        for name, value, _ in sig.iter_args(args):
+          if value is not None:
+            defaults[name] = value
+        sigs.append(sig._replace(defaults=defaults))
+    return sigs
 
   def call_slot(self, node: cfg.CFGNode, *args, **kwargs):
     return function.call_function(

pytype/tests/test_attr2.py

@@ -226,9 +226,40 @@ def f(x: int) -> str:
       @attr.s
       class Foo:
         x = attr.ib(converter=functools.partial(f))
-      # We don't yet infer the right type for Foo.x in this case, but we at
-      # least want to check that constructing a Foo doesn't generate errors.
-      Foo(x=0)
+      foo = Foo(x=0)
+      assert_type(foo.x, str)
+    """)
+
+  def test_partial_overloaded_as_converter(self):
+    self.Check("""
+      import attr
+      import functools
+      from typing import overload
+      @overload
+      def f(x: int, y: int) -> int:
+        return ''
+      @overload
+      def f(x: str, y: int) -> str:
+        return ''
+      @attr.s
+      class Foo:
+        x = attr.ib(converter=functools.partial(f, 42))
+      foo = Foo(x=0)
+      assert_type(foo.x, int)
+    """)
+
+  def test_partial_class_as_converter(self):
+    self.Check("""
+      import attr
+      import functools
+      class C:
+        def __init__(self, x: int, y: int) -> None:
+          self.x = x
+      @attr.s
+      class Foo:
+        x = attr.ib(converter=functools.partial(C, 42))
+      foo = Foo(x=0)
+      assert_type(foo.x, C)
     """)