From 1b116dfbe37a4503e0541d6bd6f5dd5c815ab36d Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Sun, 23 Jun 2024 08:47:52 +0100
Subject: [PATCH] Fix explicit type for partial (#17424)

Fixes https://github.com/python/mypy/issues/17301
---
 mypy/plugins/functools.py           | 17 +++++++++++++---
 test-data/unit/check-functools.test | 31 +++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/mypy/plugins/functools.py b/mypy/plugins/functools.py
index 4f2ed6f2361d..e41afe2fde02 100644
--- a/mypy/plugins/functools.py
+++ b/mypy/plugins/functools.py
@@ -6,6 +6,7 @@
 
 import mypy.checker
 import mypy.plugin
+import mypy.semanal
 from mypy.argmap import map_actuals_to_formals
 from mypy.nodes import ARG_POS, ARG_STAR2, ArgKind, Argument, CallExpr, FuncItem, Var
 from mypy.plugins.common import add_method_to_class
@@ -24,6 +25,8 @@
 
 _ORDERING_METHODS: Final = {"__lt__", "__le__", "__gt__", "__ge__"}
 
+PARTIAL = "functools.partial"
+
 
 class _MethodInfo(NamedTuple):
     is_static: bool
@@ -142,7 +145,8 @@ def partial_new_callback(ctx: mypy.plugin.FunctionContext) -> Type:
                 else (ArgKind.ARG_NAMED_OPT if k == ArgKind.ARG_NAMED else k)
             )
             for k in fn_type.arg_kinds
-        ]
+        ],
+        ret_type=ctx.api.named_generic_type(PARTIAL, [fn_type.ret_type]),
     )
     if defaulted.line < 0:
         # Make up a line number if we don't have one
@@ -188,6 +192,13 @@ def partial_new_callback(ctx: mypy.plugin.FunctionContext) -> Type:
     bound = get_proper_type(bound)
     if not isinstance(bound, CallableType):
         return ctx.default_return_type
+    wrapped_ret_type = get_proper_type(bound.ret_type)
+    if not isinstance(wrapped_ret_type, Instance) or wrapped_ret_type.type.fullname != PARTIAL:
+        return ctx.default_return_type
+    if not mypy.semanal.refers_to_fullname(ctx.args[0][0], PARTIAL):
+        # If the first argument is partial, above call will trigger the plugin
+        # again, in between the wrapping above an unwrapping here.
+        bound = bound.copy_modified(ret_type=wrapped_ret_type.args[0])
 
     formal_to_actual = map_actuals_to_formals(
         actual_kinds=actual_arg_kinds,
@@ -237,7 +248,7 @@ def partial_new_callback(ctx: mypy.plugin.FunctionContext) -> Type:
         ret_type=ret_type,
     )
 
-    ret = ctx.api.named_generic_type("functools.partial", [ret_type])
+    ret = ctx.api.named_generic_type(PARTIAL, [ret_type])
     ret = ret.copy_with_extra_attr("__mypy_partial", partially_applied)
     return ret
 
@@ -247,7 +258,7 @@ def partial_call_callback(ctx: mypy.plugin.MethodContext) -> Type:
     if (
         not isinstance(ctx.api, mypy.checker.TypeChecker)  # use internals
         or not isinstance(ctx.type, Instance)
-        or ctx.type.type.fullname != "functools.partial"
+        or ctx.type.type.fullname != PARTIAL
         or not ctx.type.extra_attrs
         or "__mypy_partial" not in ctx.type.extra_attrs.attrs
     ):
diff --git a/test-data/unit/check-functools.test b/test-data/unit/check-functools.test
index 79ae962a73e0..997f5bc70c7d 100644
--- a/test-data/unit/check-functools.test
+++ b/test-data/unit/check-functools.test
@@ -347,6 +347,37 @@ reveal_type(functools.partial(fn3, 2)())  # E: "str" not callable \
                                           # E: Argument 1 to "partial" has incompatible type "Union[Callable[[int], int], str]"; expected "Callable[..., int]"
 [builtins fixtures/tuple.pyi]
 
+[case testFunctoolsPartialExplicitType]
+from functools import partial
+from typing import Type, TypeVar, Callable
+
+T = TypeVar("T")
+def generic(string: str, integer: int, resulting_type: Type[T]) -> T: ...
+
+p: partial[str] = partial(generic, resulting_type=str)
+q: partial[bool] = partial(generic, resulting_type=str)  # E: Argument "resulting_type" to "generic" has incompatible type "Type[str]"; expected "Type[bool]"
+
+pc: Callable[..., str] = partial(generic, resulting_type=str)
+qc: Callable[..., bool] = partial(generic, resulting_type=str)  # E: Incompatible types in assignment (expression has type "partial[str]", variable has type "Callable[..., bool]") \
+                                                                # N: "partial[str].__call__" has type "Callable[[VarArg(Any), KwArg(Any)], str]"
+[builtins fixtures/tuple.pyi]
+
+[case testFunctoolsPartialNestedPartial]
+from functools import partial
+from typing import Any
+
+def foo(x: int) -> int: ...
+p = partial(partial, foo)
+reveal_type(p()(1))  # N: Revealed type is "builtins.int"
+p()("no")  # E: Argument 1 to "foo" has incompatible type "str"; expected "int"
+
+q = partial(partial, partial, foo)
+q()()("no")  # E: Argument 1 to "foo" has incompatible type "str"; expected "int"
+
+r = partial(partial, foo, 1)
+reveal_type(r()())  # N: Revealed type is "builtins.int"
+[builtins fixtures/tuple.pyi]
+
 [case testFunctoolsPartialTypeObject]
 import functools
 from typing import Type, Generic, TypeVar