From 9be2dea2a3f6c807a153fbf9b70a86b6fcd544ae Mon Sep 17 00:00:00 2001 From: Sergei Lebedev Date: Fri, 17 Oct 2025 08:39:28 -0700 Subject: [PATCH] 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 --- pytype/abstract/function.py | 16 ++++++++++++ pytype/overlays/functools_overlay.py | 24 ++++++++++++------ pytype/tests/test_attr2.py | 37 +++++++++++++++++++++++++--- 3 files changed, 66 insertions(+), 11 deletions(-) diff --git a/pytype/abstract/function.py b/pytype/abstract/function.py index 70422ce07..227ab4835 100644 --- a/pytype/abstract/function.py +++ b/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: diff --git a/pytype/overlays/functools_overlay.py b/pytype/overlays/functools_overlay.py index 444aec286..2fffd06cd 100644 --- a/pytype/overlays/functools_overlay.py +++ b/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( diff --git a/pytype/tests/test_attr2.py b/pytype/tests/test_attr2.py index e39305542..942cf5533 100644 --- a/pytype/tests/test_attr2.py +++ b/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) """)