From 1a074b6b2434ae2f7bbf5f12d3b0db208136dca8 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Mon, 14 Oct 2024 02:50:03 +0100
Subject: [PATCH] Better handling of generic functions in partial plugin
 (#17925)

Fixes https://github.com/python/mypy/issues/17411

The fix is that we remove type variables that can never be inferred from
the initial `check_call()` call. Actual diff is tiny, I just moved a
bunch of code, since I need formal to actual mapping sooner now.
---
 mypy/plugins/functools.py           | 62 ++++++++++++++++++-----------
 test-data/unit/check-functools.test | 31 ++++++++++++++-
 2 files changed, 68 insertions(+), 25 deletions(-)

diff --git a/mypy/plugins/functools.py b/mypy/plugins/functools.py
index f09ea88f7162..4dfeb752b5d2 100644
--- a/mypy/plugins/functools.py
+++ b/mypy/plugins/functools.py
@@ -10,6 +10,7 @@
 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
+from mypy.typeops import get_all_type_vars
 from mypy.types import (
     AnyType,
     CallableType,
@@ -17,6 +18,7 @@
     Overloaded,
     Type,
     TypeOfAny,
+    TypeVarType,
     UnboundType,
     UnionType,
     get_proper_type,
@@ -164,21 +166,6 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) -
         ctx.api.type_context[-1] = None
         wrapped_return = False
 
-    defaulted = fn_type.copy_modified(
-        arg_kinds=[
-            (
-                ArgKind.ARG_OPT
-                if k == ArgKind.ARG_POS
-                else (ArgKind.ARG_NAMED_OPT if k == ArgKind.ARG_NAMED else k)
-            )
-            for k in fn_type.arg_kinds
-        ],
-        ret_type=ret_type,
-    )
-    if defaulted.line < 0:
-        # Make up a line number if we don't have one
-        defaulted.set_line(ctx.default_return_type)
-
     # Flatten actual to formal mapping, since this is what check_call() expects.
     actual_args = []
     actual_arg_kinds = []
@@ -199,6 +186,43 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) -
             actual_arg_names.append(ctx.arg_names[i][j])
             actual_types.append(ctx.arg_types[i][j])
 
+    formal_to_actual = map_actuals_to_formals(
+        actual_kinds=actual_arg_kinds,
+        actual_names=actual_arg_names,
+        formal_kinds=fn_type.arg_kinds,
+        formal_names=fn_type.arg_names,
+        actual_arg_type=lambda i: actual_types[i],
+    )
+
+    # We need to remove any type variables that appear only in formals that have
+    # no actuals, to avoid eagerly binding them in check_call() below.
+    can_infer_ids = set()
+    for i, arg_type in enumerate(fn_type.arg_types):
+        if not formal_to_actual[i]:
+            continue
+        can_infer_ids.update({tv.id for tv in get_all_type_vars(arg_type)})
+
+    defaulted = fn_type.copy_modified(
+        arg_kinds=[
+            (
+                ArgKind.ARG_OPT
+                if k == ArgKind.ARG_POS
+                else (ArgKind.ARG_NAMED_OPT if k == ArgKind.ARG_NAMED else k)
+            )
+            for k in fn_type.arg_kinds
+        ],
+        ret_type=ret_type,
+        variables=[
+            tv
+            for tv in fn_type.variables
+            # Keep TypeVarTuple/ParamSpec to avoid spurious errors on empty args.
+            if tv.id in can_infer_ids or not isinstance(tv, TypeVarType)
+        ],
+    )
+    if defaulted.line < 0:
+        # Make up a line number if we don't have one
+        defaulted.set_line(ctx.default_return_type)
+
     # Create a valid context for various ad-hoc inspections in check_call().
     call_expr = CallExpr(
         callee=ctx.args[0][0],
@@ -231,14 +255,6 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) -
             return ctx.default_return_type
         bound = bound.copy_modified(ret_type=ret_type.args[0])
 
-    formal_to_actual = map_actuals_to_formals(
-        actual_kinds=actual_arg_kinds,
-        actual_names=actual_arg_names,
-        formal_kinds=fn_type.arg_kinds,
-        formal_names=fn_type.arg_names,
-        actual_arg_type=lambda i: actual_types[i],
-    )
-
     partial_kinds = []
     partial_types = []
     partial_names = []
diff --git a/test-data/unit/check-functools.test b/test-data/unit/check-functools.test
index bee30931a92b..ea98a902d14b 100644
--- a/test-data/unit/check-functools.test
+++ b/test-data/unit/check-functools.test
@@ -575,7 +575,6 @@ def bar(f: S) -> S:
     return f
 [builtins fixtures/primitives.pyi]
 
-
 [case testFunctoolsPartialAbstractType]
 # flags: --python-version 3.9
 from abc import ABC, abstractmethod
@@ -597,7 +596,6 @@ def f2() -> None:
     partial_cls()  # E: Cannot instantiate abstract class "A" with abstract attribute "method"
 [builtins fixtures/tuple.pyi]
 
-
 [case testFunctoolsPartialSelfType]
 from functools import partial
 from typing_extensions import Self
@@ -610,3 +608,32 @@ class A:
         factory = partial(cls, ts=0)
         return factory(msg=msg)
 [builtins fixtures/tuple.pyi]
+
+[case testFunctoolsPartialTypeVarValues]
+from functools import partial
+from typing import TypeVar
+
+T = TypeVar("T", int, str)
+
+def f(x: int, y: T) -> T:
+    return y
+
+def g(x: T, y: int) -> T:
+    return x
+
+def h(x: T, y: T) -> T:
+    return x
+
+fp = partial(f, 1)
+reveal_type(fp(1))  # N: Revealed type is "builtins.int"
+reveal_type(fp("a"))  # N: Revealed type is "builtins.str"
+fp(object())  # E: Value of type variable "T" of "f" cannot be "object"
+
+gp = partial(g, 1)
+reveal_type(gp(1))  # N: Revealed type is "builtins.int"
+gp("a")  # E: Argument 1 to "g" has incompatible type "str"; expected "int"
+
+hp = partial(h, 1)
+reveal_type(hp(1))  # N: Revealed type is "builtins.int"
+hp("a")  # E: Argument 1 to "h" has incompatible type "str"; expected "int"
+[builtins fixtures/tuple.pyi]