diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts index dfb3a548448b..c7fb9cbd1fd9 100644 --- a/packages/pyright-internal/src/analyzer/typeEvaluator.ts +++ b/packages/pyright-internal/src/analyzer/typeEvaluator.ts @@ -611,6 +611,10 @@ interface CodeFlowAnalyzerCacheEntry { codeFlowAnalyzer: CodeFlowAnalyzer; } +interface FunctionRecursionInfo { + callerNode: ExpressionNode | undefined; +} + type LogWrapper = any>(func: T) => (...args: Parameters) => ReturnType; interface SuppressedNodeStackEntry { @@ -629,7 +633,7 @@ export function createTypeEvaluator( const suppressedNodeStack: SuppressedNodeStackEntry[] = []; const assignClassToSelfStack: AssignClassToSelfInfo[] = []; - let functionRecursionMap = new Set(); + let functionRecursionMap = new Map(); let codeFlowAnalyzerCache = new Map(); let typeCache = new Map(); let effectiveTypeCache = new Map>(); @@ -669,7 +673,7 @@ export function createTypeEvaluator( // circular references in complex data structures, so it fails // to clean up the objects if we don't help it out. function disposeEvaluator() { - functionRecursionMap = new Set(); + functionRecursionMap = new Map(); codeFlowAnalyzerCache = new Map(); typeCache = new Map(); effectiveTypeCache = new Map>(); @@ -19223,7 +19227,11 @@ export function createTypeEvaluator( return awaitableReturnType; } - function inferFunctionReturnType(node: FunctionNode, isAbstract: boolean): TypeResult | undefined { + function inferFunctionReturnType( + node: FunctionNode, + isAbstract: boolean, + callerNode: ExpressionNode | undefined + ): TypeResult | undefined { const returnAnnotation = node.d.returnAnnotation || node.d.funcAnnotationComment?.d.returnAnnotation; // This shouldn't be called if there is a declared return type, but it @@ -19242,11 +19250,17 @@ export function createTypeEvaluator( return { type: inferredReturnType, isIncomplete }; } - if (functionRecursionMap.has(node.id) || functionRecursionMap.size >= maxInferFunctionReturnRecursionCount) { + const recursionEntry = functionRecursionMap.get(node.id) ?? []; + + if (functionRecursionMap.size >= maxInferFunctionReturnRecursionCount) { + inferredReturnType = UnknownType.create(); + isIncomplete = true; + } else if (recursionEntry.some((entry) => entry.callerNode === callerNode)) { inferredReturnType = UnknownType.create(); isIncomplete = true; } else { - functionRecursionMap.add(node.id); + recursionEntry.push({ callerNode }); + functionRecursionMap.set(node.id, recursionEntry); try { let functionDecl: FunctionDeclaration | undefined; @@ -19433,7 +19447,10 @@ export function createTypeEvaluator( } throw err; } finally { - functionRecursionMap.delete(node.id); + recursionEntry.pop(); + if (recursionEntry.length === 0) { + functionRecursionMap.delete(node.id); + } } } @@ -23086,7 +23103,8 @@ export function createTypeEvaluator( disableSpeculativeMode(() => { returnTypeResult = inferFunctionReturnType( functionNode, - FunctionType.isAbstractMethod(type) + FunctionType.isAbstractMethod(type), + callSiteInfo?.errorNode ); }); @@ -23295,7 +23313,8 @@ export function createTypeEvaluator( } else { contextualReturnType = inferFunctionReturnType( functionNode, - FunctionType.isAbstractMethod(type) + FunctionType.isAbstractMethod(type), + callSiteInfo?.errorNode )?.type; } }