Added heuristic to constraint solver logic to better handle the case …

…where a parameter is annotated with a union of multiple "bare" TypeVars (like `S | T`). In this case, it's not clear whether the corresponding argument type should constrain `S` or `T` or both. We now defer the constraint during the first pass of arg/param type validation so additional references to `S` or `T` can help establish the appropriate constraint. This addresses #5768. (#5772)

Co-authored-by: Eric Traut <[email protected]>

packages/pyright-internal/src/analyzer/typeEvaluator.ts

@@ -11570,14 +11570,20 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
         let isCompatible = true;
         const functionName = typeResult?.type.details.name;
         let skippedBareTypeVarExpectedType = false;
+        let skipSolveTypeVars = false;
 
         if (argParam.argument.valueExpression) {
             let expectedType: Type | undefined;
 
-            const isBareTypeVarExpectedType =
-                isTypeVar(argParam.paramType) && argParam.paramType.scopeId === typeResult?.type.details.typeVarScopeId;
+            // Is the expected type a "bare" in-scope TypeVar or a union of bare in-scope TypeVars?
+            let isExpectedTypeBareTypeVar = true;
+            doForEachSubtype(argParam.paramType, (subtype) => {
+                if (!isTypeVar(subtype) || subtype.scopeId !== typeResult?.type.details.typeVarScopeId) {
+                    isExpectedTypeBareTypeVar = false;
+                }
+            });
 
-            if (!options.skipBareTypeVarExpectedType || !isBareTypeVarExpectedType) {
+            if (!options.skipBareTypeVarExpectedType || !isExpectedTypeBareTypeVar) {
                 expectedType = argParam.paramType;
 
                 // If the parameter type is a function with a ParamSpec, don't apply
@@ -11595,6 +11601,14 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
                 }
             } else {
                 skippedBareTypeVarExpectedType = true;
+
+                // If the expected type is a union of bare TypeVars, it's not clear which of the two
+                // (or both) should be constrained. We'll skip any attempt to solve the TypeVars during
+                // this pass and hope that subsequent arg assignments will help us establish the correct
+                // constraints.
+                if (isUnion(argParam.paramType)) {
+                    skipSolveTypeVars = true;
+                }
             }
 
             // If the expected type is unknown, don't use an expected type. Instead,
@@ -11753,7 +11767,16 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
             }
         }
 
-        if (!assignType(argParam.paramType, argType, diag.createAddendum(), typeVarContext)) {
+        if (
+            !assignType(
+                argParam.paramType,
+                argType,
+                diag.createAddendum(),
+                typeVarContext,
+                /* srcTypeVarContext */ undefined,
+                skipSolveTypeVars ? AssignTypeFlags.SkipSolveTypeVars : undefined
+            )
+        ) {
             // Mismatching parameter types are common in untyped code; don't bother spending time
             // printing types if the diagnostic is disabled.
             const fileInfo = AnalyzerNodeInfo.getFileInfo(argParam.errorNode);

packages/pyright-internal/src/tests/samples/solver29.py

@@ -0,0 +1,36 @@
+# This sample tests the case where a union of two "bare" TypeVars are
+# used in the annotation of a function parameter.
+
+from typing import Any, TypeVar, Callable
+
+S = TypeVar("S")
+T = TypeVar("T")
+
+
+def accepts_bool(b: bool) -> None:
+    ...
+
+
+def accepts_int(i: int) -> None:
+    ...
+
+
+def func1(x: S | T, l2: Callable[[S], Any], l3: Callable[[T], Any]) -> tuple[S, T]:
+    ...
+
+
+def func2(x: T | S, l2: Callable[[S], Any], l3: Callable[[T], Any]) -> tuple[S, T]:
+    ...
+
+
+x1 = func1(0, accepts_int, accepts_bool)
+reveal_type(x1, expected_text="tuple[int, bool]")
+
+x2 = func1(True, accepts_int, accepts_bool)
+reveal_type(x2, expected_text="tuple[int, bool]")
+
+x3 = func1(0, accepts_int, accepts_bool)
+reveal_type(x3, expected_text="tuple[int, bool]")
+
+x4 = func1(True, accepts_int, accepts_bool)
+reveal_type(x4, expected_text="tuple[int, bool]")

packages/pyright-internal/src/tests/typeEvaluator2.test.ts

@@ -671,6 +671,12 @@ test('Solver28', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('Solver29', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['solver29.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('SolverScoring1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['solverScoring1.py']);