diff --git a/packages/pyright-internal/src/analyzer/binder.ts b/packages/pyright-internal/src/analyzer/binder.ts index 07ec96ecdbe7..77ea07d808f8 100644 --- a/packages/pyright-internal/src/analyzer/binder.ts +++ b/packages/pyright-internal/src/analyzer/binder.ts @@ -90,6 +90,7 @@ import { FlowExhaustedMatch, FlowFlags, FlowLabel, + FlowLoopLabel, FlowNarrowForPattern, FlowNode, FlowPostContextManagerLabel, @@ -2314,10 +2315,11 @@ export class Binder extends ParseTreeWalker { } private _createLoopLabel() { - const flowNode: FlowLabel = { + const flowNode: FlowLoopLabel = { flags: FlowFlags.LoopLabel, id: getUniqueFlowNodeId(), antecedents: [], + expressions: undefined, }; return flowNode; } @@ -2889,14 +2891,25 @@ export class Binder extends ParseTreeWalker { } } - private _bindLoopStatement(preLoopLabel: FlowLabel, postLoopLabel: FlowLabel, callback: () => void) { + private _bindLoopStatement(preLoopLabel: FlowLoopLabel, postLoopLabel: FlowLabel, callback: () => void) { const savedContinueTarget = this._currentContinueTarget; const savedBreakTarget = this._currentBreakTarget; + const savedExpressions = this._currentScopeCodeFlowExpressions; this._currentContinueTarget = preLoopLabel; this._currentBreakTarget = postLoopLabel; + this._currentScopeCodeFlowExpressions = new Set(); callback(); + preLoopLabel.expressions = this._currentScopeCodeFlowExpressions; + + if (savedExpressions) { + this._currentScopeCodeFlowExpressions.forEach((value) => { + savedExpressions.add(value); + }); + } + + this._currentScopeCodeFlowExpressions = savedExpressions; this._currentContinueTarget = savedContinueTarget; this._currentBreakTarget = savedBreakTarget; } diff --git a/packages/pyright-internal/src/analyzer/codeFlow.ts b/packages/pyright-internal/src/analyzer/codeFlow.ts index 0969ca4b828d..c528359a5357 100644 --- a/packages/pyright-internal/src/analyzer/codeFlow.ts +++ b/packages/pyright-internal/src/analyzer/codeFlow.ts @@ -71,6 +71,14 @@ export interface FlowLabel extends FlowNode { antecedents: FlowNode[]; } +export interface FlowLoopLabel extends FlowLabel { + // Set of all expressions that require code flow analysis + // through the loop to determine their types. If an expression + // is not within this map, loop analysis can be skipped and + // determined from the first antecedent only. + expressions: Set | undefined; +} + // FlowAssignment represents a node that assigns a value. export interface FlowAssignment extends FlowNode { node: CodeFlowReferenceExpressionNode; diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts index 7d5322f1c19f..b27e12c017ad 100644 --- a/packages/pyright-internal/src/analyzer/typeEvaluator.ts +++ b/packages/pyright-internal/src/analyzer/typeEvaluator.ts @@ -93,6 +93,7 @@ import { FlowExhaustedMatch, FlowFlags, FlowLabel, + FlowLoopLabel, FlowNarrowForPattern, FlowNode, FlowPostContextManagerLabel, @@ -17576,7 +17577,18 @@ export function createTypeEvaluator( } if (curFlowNode.flags & FlowFlags.LoopLabel) { - const loopNode = curFlowNode as FlowLabel; + const loopNode = curFlowNode as FlowLoopLabel; + + // Is the current symbol modified in any way within the + // loop? If not, we can skip all processing within the loop + // and assume that the type comes from the first antecedent, + // which feeds the loop. + if (referenceKey) { + if (!loopNode.expressions || !loopNode.expressions.has(referenceKey)) { + curFlowNode = loopNode.antecedents[0]; + continue; + } + } let firstWasIncomplete = false; let isFirstTimeInLoop = false;