@zzzen/pyright-internal 1.2.0-dev.20261011 → 1.2.0-dev.20261018

package/dist/analyzer/typeEvaluator.js

@@ -80,6 +80,7 @@ const namedTuples_1 = require("./namedTuples");
 const operations_1 = require("./operations");
 const parameterUtils_1 = require("./parameterUtils");
 const ParseTreeUtils = __importStar(require("./parseTreeUtils"));
+const parseTreeWalker_1 = require("./parseTreeWalker");
 const patternMatching_1 = require("./patternMatching");
 const properties_1 = require("./properties");
 const protocols_1 = require("./protocols");
@@ -203,6 +204,7 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
     let functionRecursionMap = new Map();
     let codeFlowAnalyzerCache = new Map();
     let typeCache = new Map();
+    let contextualLambdaTypeCache;
     let typeFormTypeCache = new Map();
     let effectiveTypeCache = new Map();
     let exceptionTargetSymbolCache = new Map();
@@ -399,6 +401,55 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
     function setTypeResultForNode(node, typeResult, flags = 0 /* EvalFlags.None */) {
         writeTypeCache(node, typeResult, flags);
     }
+    function captureContextualLambdaTypes(node, typeResult) {
+        const cache = contextualLambdaTypeCache;
+        if (!cache || !cache.lambdaNodes.some((root) => ParseTreeUtils.isNodeContainedWithin(node, root))) {
+            return;
+        }
+        captureLambdaTypeCacheEntries(node, cache, typeResult);
+    }
+    function captureLambdaTypeCacheEntries(node, cache, typeResult) {
+        checkForCancellation();
+        const entry = typeResult ? { typeResult, flags: 0 /* EvalFlags.None */, incompleteGenCount } : readTypeCacheEntry(node);
+        if (entry && !entry.typeResult.isIncomplete) {
+            captureLambdaTypeCacheEntry(node, cache, entry);
+        }
+        if (typeResult) {
+            return;
+        }
+        for (const child of (0, parseTreeWalker_1.getChildNodes)(node)) {
+            // Nested lambdas capture their own contextual types before any
+            // non-contextual fallback can replace them.
+            if (child && child.nodeType !== 33 /* ParseNodeType.Lambda */) {
+                captureLambdaTypeCacheEntries(child, cache);
+            }
+        }
+    }
+    function captureLambdaTypeCacheEntry(node, cache, entry) {
+        const nodeEntries = cache.entries.get(node) ?? new Map();
+        // Generic inference can revisit a lambda within the same callable branch.
+        nodeEntries.set(cache.branchIndex, entry);
+        cache.entries.set(node, nodeEntries);
+    }
+    function combineContextualLambdaTypes(cache) {
+        for (const [node, branchEntries] of cache.entries) {
+            checkForCancellation();
+            const nodeEntries = Array.from(branchEntries.values());
+            const lastEntry = nodeEntries[nodeEntries.length - 1];
+            writeTypeCache(node, {
+                ...lastEntry.typeResult,
+                type: (0, types_1.combineTypes)(nodeEntries.map((entry) => entry.typeResult.type)),
+                typeErrors: nodeEntries.some((entry) => entry.typeResult.typeErrors),
+            }, readTypeCacheEntry(node)?.flags ?? lastEntry.flags);
+            const parentCache = contextualLambdaTypeCache;
+            const combinedEntry = readTypeCacheEntry(node);
+            if (parentCache &&
+                combinedEntry &&
+                parentCache.lambdaNodes.some((root) => ParseTreeUtils.isNodeContainedWithin(node, root))) {
+                captureLambdaTypeCacheEntry(node, parentCache, combinedEntry);
+            }
+        }
+    }
     function getMaxCodeComplexity() {
         return maxCodeComplexity;
     }
@@ -6100,12 +6151,22 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
                     }
                 }
             }
-            else if ((0, types_1.isClassInstance)(valueType) && types_1.ClassType.isBuiltIn(valueType, 'slice')) {
+            else if (usage.method === 'get' &&
+                index0Expr.nodeType === 46 /* ParseNodeType.Slice */ &&
+                (0, types_1.isClassInstance)(valueType) &&
+                types_1.ClassType.isBuiltIn(valueType, 'slice')) {
                 const tupleType = (0, typeUtils_1.getSpecializedTupleType)(baseType);
-                if (tupleType && index0Expr.nodeType === 46 /* ParseNodeType.Slice */) {
-                    const slicedTupleType = (0, tuples_1.getSlicedTupleType)(evaluatorInterface, tupleType, index0Expr);
-                    if (slicedTupleType) {
-                        return { type: slicedTupleType };
+                if (tupleType) {
+                    // Skip this if a subclass overrides the tuple's __getitem__ method.
+                    const getItemMember = (0, typeUtils_1.lookUpObjectMember)(baseType, '__getitem__', 16 /* MemberAccessFlags.SkipInstanceMembers */);
+                    const isTupleGetItem = !getItemMember ||
+                        ((0, types_1.isInstantiableClass)(getItemMember.classType) &&
+                            types_1.ClassType.isBuiltIn(getItemMember.classType, 'tuple'));
+                    if (isTupleGetItem) {
+                        const slicedTupleType = (0, tuples_1.getSlicedTupleType)(evaluatorInterface, tupleType, index0Expr);
+                        if (slicedTupleType) {
+                            return { type: slicedTupleType };
+                        }
                     }
                 }
             }
@@ -7218,7 +7279,7 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
     // there will be only one argument list in expandedArgTypes, and all entries
     // (one for each argument) will be undefined. On subsequent calls, this
     // list will grow to include union expansions.
-    function validateOverloadsWithExpandedTypes(errorNode, expandedArgTypes, argParamMatches, constraints, skipUnknownArgCheck, inferenceContext) {
+    function validateOverloadsWithExpandedTypes(errorNode, expandedArgTypes, argParamMatches, constraints, skipUnknownArgCheck, inferenceContext, skipFinalValidation = false) {
         const returnTypes = [];
         const specializedInitSelfTypes = [];
         let matchedOverloads = [];
@@ -7261,17 +7322,14 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
                 if (!callResult.argumentErrors && callResult.returnType) {
                     if (evaluatorOptions.experimentalOverloadResult &&
                         !callResult.isTypeIncomplete &&
-                        !overload.priv.boundToType &&
-                        !overload.shared.methodClass &&
-                        !overload.shared.declaration?.isMethod &&
                         errorNode.nodeType === 9 /* ParseNodeType.Call */ &&
-                        errorNode.d.leftExpr.nodeType === 38 /* ParseNodeType.Name */ &&
-                        errorNode.parent?.nodeType === 3 /* ParseNodeType.Assignment */ &&
-                        errorNode.parent.d.rightExpr === errorNode &&
+                        (0, overloadResultSelector_1.isAutomaticOverloadSeedSyntax)(errorNode) &&
                         expandedArgTypes.length === 1 &&
-                        !hasArgTypeOverride &&
-                        callResult.argResults) {
-                        evaluatorOptions.experimentalOverloadResult.considerActivation(errorNode, callResult.argResults);
+                        !hasArgTypeOverride) {
+                        const receiver = overload.priv.boundToType;
+                        evaluatorOptions.experimentalOverloadResult.considerActivation(errorNode, receiver && (0, types_1.isClassInstance)(receiver)
+                            ? [{ argType: receiver }, ...(callResult.argResults ?? [])]
+                            : callResult.argResults ?? []);
                     }
                     overloadsUsedForCall.push(overload);
                     matchedOverload = overload;
@@ -7390,6 +7448,16 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
         if (constraints && isDefinitiveMatchFound) {
             constraints.copyFromClone(matchedOverloads[matchedOverloads.length - 1].constraints);
         }
+        // A speculative refinement does not need to populate the type cache;
+        // the original call has already performed the normal final validation.
+        if (skipFinalValidation) {
+            return {
+                argumentErrors: false,
+                returnType: (0, types_1.combineTypes)(returnTypes),
+                isTypeIncomplete,
+                overloadsUsedForCall,
+            };
+        }
         // And run through the first expanded argument list one more time to
         // populate the type cache, without applying a different branch's constraints.
         const finalConstraints = expandedArgTypes.length > 1
@@ -7518,10 +7586,13 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
                 name: arg.d.name,
             }));
             const sources = node.d.args.map((arg) => getTypeOfExpression(arg.d.valueExpr));
-            if (sources.some((s) => s.isIncomplete || s.typeErrors || (0, types_1.isAny)(s.type) || !(0, overloadResultSelector_1.isFixedOverloadArgumentShape)(s.type, budget))) {
+            const receiver = node.d.leftExpr.nodeType === 35 /* ParseNodeType.MemberAccess */
+                ? getTypeOfExpression(node.d.leftExpr.d.leftExpr)
+                : undefined;
+            if ([...sources, ...(receiver ? [receiver] : [])].some((s) => s.isIncomplete || s.typeErrors || (0, types_1.isAny)(s.type) || !(0, overloadResultSelector_1.isFixedOverloadArgumentShape)(s.type, budget))) {
                 return decline('unsupported-argument-shape');
             }
-            if (!sources.some((s) => (0, overloadResultSelector_1.hasNestedOverloadAny)(s.type))) {
+            if (![...sources, ...(receiver ? [receiver] : [])].some((s) => (0, overloadResultSelector_1.hasNestedOverloadAny)(s.type))) {
                 return decline('no-invariant-any');
             }
             const returns = [];
@@ -7531,10 +7602,22 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
                 const declaration = overload.shared.declaration;
                 const declaredReturnType = overload.shared.declaredReturnType;
                 const returnType = types_1.FunctionType.getEffectiveReturnType(overload);
+                const receiverParam = overload.priv.strippedFirstParamType;
+                const supportedBinding = receiver
+                    ? declaration?.isMethod &&
+                        types_1.FunctionType.isInstanceMethod(overload) &&
+                        overload.priv.boundToType &&
+                        (0, types_1.isClassInstance)(overload.priv.boundToType) &&
+                        receiverParam &&
+                        (0, types_1.isClassInstance)(receiverParam) &&
+                        (0, types_1.isClassInstance)(receiver.type) &&
+                        types_1.ClassType.isSameGenericClass(receiver.type, receiverParam) &&
+                        overload.shared.methodClass &&
+                        types_1.ClassType.isSameGenericClass(receiver.type, types_1.ClassType.cloneAsInstance(overload.shared.methodClass)) &&
+                        (0, overloadResultSelector_1.isFixedOverloadArgumentShape)(receiverParam, budget)
+                    : !declaration?.isMethod && !overload.shared.methodClass && !overload.priv.boundToType;
                 if (!declaration ||
-                    declaration.isMethod ||
-                    overload.shared.methodClass ||
-                    overload.priv.boundToType ||
+                    !supportedBinding ||
                     overload.shared.typeParams.length ||
                     !declaredReturnType ||
                     !returnType ||
@@ -7576,6 +7659,9 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
                     continue;
                 }
                 const proof = mapping.argParams.map((p, i) => (0, overloadResultSelector_1.getFixedOverloadCoverage)(sources[i].type, p.paramType, budget));
+                if (receiver && receiverParam) {
+                    proof.unshift((0, overloadResultSelector_1.getFixedOverloadCoverage)(receiver.type, receiverParam, budget));
+                }
                 const covered = proof.includes('unsupported')
                     ? 'unsupported'
                     : proof.includes('not-covered')
@@ -7600,11 +7686,12 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
                 : decline('equivalent-or-single-return');
         });
     }
-    function validateOverloadedArgTypes(errorNode, argList, typeResult, constraints, skipUnknownArgCheck, inferenceContext) {
+    function validateOverloadedArgTypes(errorNode, argList, typeResult, constraints, skipUnknownArgCheck, inferenceContext, candidateOverloads) {
         const filteredMatchResults = [];
         let contextFreeArgTypes;
         let isTypeIncomplete = !!typeResult.isIncomplete;
         const type = typeResult.type;
+        const overloads = candidateOverloads ?? types_1.OverloadedType.getOverloads(type);
         const speculativeNode = getSpeculativeNodeForCall(errorNode);
         // Start by evaluating the types of the arguments without any expected
         // type. Also, filter the list of overloads based on the number of
@@ -7613,7 +7700,7 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
         // cache or record any diagnostics at this stage.
         useSpeculativeMode(speculativeNode, () => {
             let overloadIndex = 0;
-            types_1.OverloadedType.getOverloads(type).forEach((overload) => {
+            overloads.forEach((overload) => {
                 // Consider only the functions that have the @overload decorator,
                 // not the final function that omits the overload. This is the
                 // intended behavior according to PEP 484.
@@ -7630,7 +7717,6 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
             // Skip the error message if we're in speculative mode because it's very
             // expensive, and we're going to suppress the diagnostic anyway.
             if (!canSkipDiagnosticForNode(errorNode)) {
-                const overloads = types_1.OverloadedType.getOverloads(type);
                 const functionName = overloads.length > 0 && overloads[0].shared.name
                     ? overloads[0].shared.name
                     : '<anonymous function>';
@@ -7685,7 +7771,8 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
         }
         let expandedArgTypes = [argList.map((arg) => undefined)];
         while (true) {
-            const callResult = validateOverloadsWithExpandedTypes(errorNode, expandedArgTypes, filteredMatchResults, constraints, skipUnknownArgCheck, inferenceContext);
+            const callResult = validateOverloadsWithExpandedTypes(errorNode, expandedArgTypes, filteredMatchResults, constraints, skipUnknownArgCheck, inferenceContext, 
+            /* skipFinalValidation */ candidateOverloads !== undefined);
             if (callResult.isTypeIncomplete) {
                 isTypeIncomplete = true;
             }
@@ -7820,24 +7907,51 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
             }), exprNode);
             return { returnType: types_1.UnknownType.create(), argumentErrors: true, overloadsUsedForCall };
         }
-        let returnType = mapSubtypesExpandTypeVars(callTypeResult.type, { sortSubtypes: true }, (expandedSubtype, unexpandedSubtype, isLastIteration) => {
-            return useSpeculativeMode(isLastIteration ? undefined : getSpeculativeNodeForCall(errorNode), () => {
-                const callResult = validateCallArgsForSubtype(errorNode, argList, expandedSubtype, unexpandedSubtype, !!callTypeResult.isIncomplete, constraints, skipUnknownArgCheck, inferenceContext, recursionCount);
-                if (callResult.argumentErrors) {
-                    argumentErrors = true;
-                }
-                if (callResult.isTypeIncomplete) {
-                    isTypeIncomplete = true;
-                }
-                if (callResult.overloadsUsedForCall) {
-                    (0, collectionUtils_1.appendArray)(overloadsUsedForCall, callResult.overloadsUsedForCall);
+        const concreteCallType = makeTopLevelTypeVarsConcrete(callTypeResult.type);
+        const lambdaNodes = (0, types_1.isUnion)(concreteCallType)
+            ? argList
+                .map((arg) => arg.valueExpression)
+                .filter((node) => node?.nodeType === 33 /* ParseNodeType.Lambda */)
+            : undefined;
+        const lambdaTypeCache = lambdaNodes?.length
+            ? { lambdaNodes, entries: new Map(), branchIndex: 0 }
+            : undefined;
+        const previousCache = contextualLambdaTypeCache;
+        let returnType;
+        try {
+            if (lambdaTypeCache) {
+                contextualLambdaTypeCache = lambdaTypeCache;
+            }
+            returnType = mapSubtypesExpandTypeVars(callTypeResult.type, { sortSubtypes: true }, (expandedSubtype, unexpandedSubtype, isLastIteration) => {
+                if (lambdaTypeCache) {
+                    lambdaTypeCache.branchIndex++;
                 }
-                specializedInitSelfType = callResult.specializedInitSelfType;
-                return callResult.returnType;
-            }, {
-                allowDiagnostics: true,
+                return useSpeculativeMode(isLastIteration ? undefined : getSpeculativeNodeForCall(errorNode), () => {
+                    const callResult = validateCallArgsForSubtype(errorNode, argList, expandedSubtype, unexpandedSubtype, !!callTypeResult.isIncomplete, constraints, skipUnknownArgCheck, inferenceContext, recursionCount);
+                    if (callResult.argumentErrors) {
+                        argumentErrors = true;
+                    }
+                    if (callResult.isTypeIncomplete) {
+                        isTypeIncomplete = true;
+                    }
+                    if (callResult.overloadsUsedForCall) {
+                        (0, collectionUtils_1.appendArray)(overloadsUsedForCall, callResult.overloadsUsedForCall);
+                    }
+                    specializedInitSelfType = callResult.specializedInitSelfType;
+                    return callResult.returnType;
+                }, {
+                    allowDiagnostics: true,
+                });
             });
-        });
+        }
+        finally {
+            contextualLambdaTypeCache = previousCache;
+        }
+        // Keep branch-specific checking and return inference separate from the
+        // combined contextual types retained for subsequent queries.
+        if (!isTypeIncomplete && lambdaTypeCache && lambdaTypeCache.branchIndex > 1) {
+            (0, cancellationUtils_1.invalidateTypeCacheIfCanceled)(() => combineContextualLambdaTypes(lambdaTypeCache));
+        }
         // If we ended up with a "Never" type because all code paths returned
         // undefined due to argument errors, transform the result into an Unknown
         // to avoid subsequent false positives.
@@ -8037,6 +8151,110 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
         }
         return { symbol, symbolName, classType, hasImplementation };
     }
+    // The synthesized "get" and "pop" methods for a TypedDict include a
+    // catch-all overload that accepts any str key. When the key argument is a
+    // union of str literals, that catch-all overload matches before union
+    // expansion is attempted, which produces an overly-wide return type. This
+    // function evaluates the call separately for each literal key subtype
+    // (speculatively) and, if every subtype is accepted, returns the union of
+    // the resulting return types.
+    function refineTypedDictMethodCallForUnionKey(errorNode, argList, overloadedType, isCallTypeIncomplete, constraints, inferenceContext) {
+        const overloads = types_1.OverloadedType.getOverloads(overloadedType);
+        const overload = overloads[0];
+        if (!overload || !types_1.FunctionType.isSynthesizedMethod(overload)) {
+            return undefined;
+        }
+        const methodName = overload.shared.name;
+        if (methodName !== 'get' && methodName !== 'pop') {
+            return undefined;
+        }
+        const isBound = (0, types_1.isMethodType)(overload);
+        const receiverType = isBound
+            ? overload.priv.strippedFirstParamType
+            : overload.shared.parameters.length > 0
+                ? types_1.FunctionType.getParamType(overload, 0)
+                : undefined;
+        if (!receiverType || !(0, types_1.isClassInstance)(receiverType) || !types_1.ClassType.isTypedDictClass(receiverType)) {
+            return undefined;
+        }
+        // Unbound generic methods have not specialized their field types from the
+        // receiver. Keep the normal result rather than refining it to Unknown.
+        if (!isBound &&
+            ((0, typeUtils_1.requiresSpecialization)(receiverType) || (0, typeUtils_1.containsAnyOrUnknown)(receiverType, /* recurse */ true))) {
+            return undefined;
+        }
+        // Handle only simple positional calls: (key) or (key, default).
+        const keyIndex = isBound ? 0 : 1;
+        if (argList.length <= keyIndex || argList.length > keyIndex + 2) {
+            return undefined;
+        }
+        if (!argList.every((arg) => !arg.name && arg.argCategory === 0 /* ArgCategory.Simple */)) {
+            return undefined;
+        }
+        const keyArg = argList[keyIndex];
+        const keyTypeResult = keyArg.typeResult ?? (keyArg.valueExpression ? getTypeOfExpression(keyArg.valueExpression) : undefined);
+        if (!keyTypeResult || !(0, types_1.isUnion)(keyTypeResult.type)) {
+            return undefined;
+        }
+        // Every subtype must be a str literal; otherwise keep the normal result.
+        const keySubtypes = [];
+        let allStrLiterals = true;
+        (0, typeUtils_1.doForEachSubtype)(keyTypeResult.type, (subtype) => {
+            if ((0, types_1.isClassInstance)(subtype) &&
+                types_1.ClassType.isBuiltIn(subtype, 'str') &&
+                typeof subtype.priv.literalValue === 'string') {
+                keySubtypes.push(subtype);
+            }
+            else {
+                allStrLiterals = false;
+            }
+        });
+        if (!allStrLiterals || keySubtypes.length < 2 || keySubtypes.length > maxSingleOverloadArgTypeExpansionCount) {
+            return undefined;
+        }
+        let isTypeIncomplete = !!keyTypeResult.isIncomplete;
+        const returnTypes = [];
+        const overloadsUsedForCall = new Set();
+        const succeeded = useSpeculativeMode(getSpeculativeNodeForCall(errorNode), () => {
+            for (const keySubtype of keySubtypes) {
+                // Avoid revalidating unrelated keys and their default parameters.
+                const candidates = overloads.filter((candidate) => {
+                    const paramType = types_1.FunctionType.getParamType(candidate, keyIndex);
+                    return ((0, types_1.isClassInstance)(paramType) &&
+                        types_1.ClassType.isBuiltIn(paramType, 'str') &&
+                        (paramType.priv.literalValue === undefined ||
+                            paramType.priv.literalValue === keySubtype.priv.literalValue));
+                });
+                if (candidates.length === 0) {
+                    return false;
+                }
+                const subtypeArgList = [...argList];
+                subtypeArgList[keyIndex] = {
+                    argCategory: 0 /* ArgCategory.Simple */,
+                    typeResult: { type: keySubtype, isIncomplete: keyTypeResult.isIncomplete },
+                };
+                const subtypeResult = validateOverloadedArgTypes(errorNode, subtypeArgList, { type: overloadedType, isIncomplete: isCallTypeIncomplete }, constraints?.clone(), 
+                /* skipUnknownArgCheck */ true, inferenceContext, candidates);
+                if (subtypeResult.isTypeIncomplete) {
+                    isTypeIncomplete = true;
+                }
+                if (subtypeResult.argumentErrors || !subtypeResult.returnType) {
+                    return false;
+                }
+                returnTypes.push(subtypeResult.returnType);
+                subtypeResult.overloadsUsedForCall?.forEach((matched) => overloadsUsedForCall.add(matched));
+            }
+            return true;
+        });
+        if (!succeeded) {
+            return undefined;
+        }
+        return {
+            returnType: (0, types_1.combineTypes)(returnTypes),
+            isTypeIncomplete,
+            overloadsUsedForCall: [...overloadsUsedForCall],
+        };
+    }
     function validateCallForOverloaded(errorNode, argList, expandedCallType, isCallTypeIncomplete, constraints, skipUnknownArgCheck, inferenceContext) {
         const overloads = types_1.OverloadedType.getOverloads(expandedCallType);
         // Handle the 'cast' call as a special case.
@@ -8049,11 +8267,19 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
         let returnType = callResult.returnType ?? types_1.UnknownType.create();
         let isTypeIncomplete = !!callResult.isTypeIncomplete;
         let argumentErrors = !!callResult.argumentErrors;
+        let overloadsUsedForCall = callResult.overloadsUsedForCall;
         if (!argumentErrors) {
+            const refinedResult = refineTypedDictMethodCallForUnionKey(errorNode, argList, expandedCallType, isCallTypeIncomplete, constraints, inferenceContext);
+            if (refinedResult?.isTypeIncomplete) {
+                isTypeIncomplete = true;
+            }
+            if (refinedResult) {
+                overloadsUsedForCall = refinedResult.overloadsUsedForCall;
+            }
             // Call the function transform logic to handle special-cased functions.
             const transformed = (0, functionTransform_1.applyFunctionTransform)(evaluatorInterface, errorNode, argList, expandedCallType, {
                 argumentErrors: !!callResult.argumentErrors,
-                returnType: callResult.returnType ?? types_1.UnknownType.create(isTypeIncomplete),
+                returnType: refinedResult?.returnType ?? callResult.returnType ?? types_1.UnknownType.create(isTypeIncomplete),
                 isTypeIncomplete,
             });
             returnType = transformed.returnType;
@@ -8068,7 +8294,7 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
             returnType,
             isTypeIncomplete,
             argumentErrors,
-            overloadsUsedForCall: callResult.overloadsUsedForCall,
+            overloadsUsedForCall,
             specializedInitSelfType: callResult.specializedInitSelfType,
         };
     }
@@ -11496,6 +11722,9 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
                         }
                     }
                 }
+                if (expectedType && !forceSpeculative && !isIncomplete) {
+                    captureContextualLambdaTypes(node);
+                }
             }, {
                 dependentType: inferenceContext?.expectedType,
                 allowDiagnostics: !forceSpeculative && !canSkipDiagnosticForNode(node) && !inferenceContext?.isTypeIncomplete,
@@ -11506,7 +11735,11 @@ function createTypeEvaluator(importLookup, evaluatorOptions, wrapWithLogger) {
             if (expectedType && !assignType(expectedType, functionType)) {
                 typeErrors = true;
             }
-            return { type: functionType, isIncomplete, typeErrors };
+            const typeResult = { type: functionType, isIncomplete, typeErrors };
+            if (expectedType && !forceSpeculative && !isIncomplete) {
+                captureContextualLambdaTypes(node, typeResult);
+            }
+            return typeResult;
         });
     }
     function getTypeOfComprehension(node, flags, inferenceContext) {