From 1a61e1c4a3aa950c9c6db1673b529f296f5705fe Mon Sep 17 00:00:00 2001 From: Jackson Cooper Date: Sat, 24 Jan 2026 04:10:13 +1100 Subject: [PATCH] Add test and fix for ternary expression narrowing issue --- .../src/analyzer/operations.ts | 50 ++++++++- .../src/languageServerBase.ts | 11 +- .../src/tests/languageServer.test.ts | 83 +++++++++----- .../src/tests/lsp/languageServerTestUtils.ts | 2 +- .../src/tests/samples/conditionalExpr1.py | 106 ++++++++++++++++++ .../src/tests/typeEvaluator5.test.ts | 5 + .../pyright-internal/src/workspaceFactory.ts | 6 +- 7 files changed, 221 insertions(+), 42 deletions(-) create mode 100644 packages/pyright-internal/src/tests/samples/conditionalExpr1.py diff --git a/packages/pyright-internal/src/analyzer/operations.ts b/packages/pyright-internal/src/analyzer/operations.ts index d4ed00b39d5c..02aea0fd4e59 100644 --- a/packages/pyright-internal/src/analyzer/operations.ts +++ b/packages/pyright-internal/src/analyzer/operations.ts @@ -765,6 +765,33 @@ export function getTypeOfUnaryOperation( return { type, isIncomplete, magicMethodDeprecationInfo: deprecatedInfo }; } +// Helper function to check if an expression is a simple name or `not ` with a literal bool type. +// We avoid narrowing for these cases because the variable could be reassigned. +function isBoolLiteralName(expr: ExpressionNode, exprType: Type, evaluator: TypeEvaluator): boolean { + // Check for simple name references + if (expr.nodeType === ParseNodeType.Name) { + return ( + isClassInstance(exprType) && + ClassType.isBuiltIn(exprType, 'bool') && + exprType.priv.literalValue !== undefined + ); + } + + // Check for `not ` expressions + if (expr.nodeType === ParseNodeType.UnaryOperation && expr.d.operator === OperatorType.Not) { + // Evaluate the inner expression's type (not the `not` expression's type). + // Note: makeTopLevelTypeVarsConcrete is needed here for the isClassInstance/ + // ClassType.isBuiltIn checks, unlike the outer call site which was removed + // because canBeTruthy/canBeFalsy do it internally. + const innerType = evaluator.makeTopLevelTypeVarsConcrete( + evaluator.getTypeOfExpression(expr.d.expr).type + ); + return isBoolLiteralName(expr.d.expr, innerType, evaluator); + } + + return false; +} + export function getTypeOfTernaryOperation( evaluator: TypeEvaluator, node: TernaryNode, @@ -778,7 +805,9 @@ export function getTypeOfTernaryOperation( return { type: UnknownType.create() }; } - evaluator.getTypeOfExpression(node.d.testExpr); + // Get the narrowed type of the test expression at this point in the code flow. + const testExprTypeResult = evaluator.getTypeOfExpression(node.d.testExpr); + const testExprType = testExprTypeResult.type; const typesToCombine: Type[] = []; let isIncomplete = false; @@ -790,7 +819,19 @@ export function getTypeOfTernaryOperation( fileInfo.definedConstants ); - if (constExprValue !== false && evaluator.isNodeReachable(node.d.ifExpr)) { + // Check if we should apply flow-sensitive narrowing. We avoid narrowing for + // simple name references with literal bool types because the variable could + // be reassigned, even though the type is a literal. This also applies to + // `not ` expressions to maintain consistency. + // Note: This guard is specific to ternary expressions. The and/or operators + // don't need this guard because they operate on already-evaluated types from + // their operands, not on types that may have been narrowed by upstream flow analysis. + const shouldApplyNarrowing = !isBoolLiteralName(node.d.testExpr, testExprType, evaluator); + + // Determine if the if-branch is reachable based on static evaluation, + // general reachability, and flow-sensitive type narrowing. + const testCanBeTruthy = shouldApplyNarrowing ? evaluator.canBeTruthy(testExprType) : true; + if (constExprValue !== false && evaluator.isNodeReachable(node.d.ifExpr) && testCanBeTruthy) { const ifType = evaluator.getTypeOfExpression(node.d.ifExpr, flags, inferenceContext); typesToCombine.push(ifType.type); if (ifType.isIncomplete) { @@ -801,7 +842,10 @@ export function getTypeOfTernaryOperation( } } - if (constExprValue !== true && evaluator.isNodeReachable(node.d.elseExpr)) { + // Determine if the else-branch is reachable based on static evaluation, + // general reachability, and flow-sensitive type narrowing. + const testCanBeFalsy = shouldApplyNarrowing ? evaluator.canBeFalsy(testExprType) : true; + if (constExprValue !== true && evaluator.isNodeReachable(node.d.elseExpr) && testCanBeFalsy) { const elseType = evaluator.getTypeOfExpression(node.d.elseExpr, flags, inferenceContext); typesToCombine.push(elseType.type); if (elseType.isIncomplete) { diff --git a/packages/pyright-internal/src/languageServerBase.ts b/packages/pyright-internal/src/languageServerBase.ts index 5945a5bb9e87..6addb98c799d 100644 --- a/packages/pyright-internal/src/languageServerBase.ts +++ b/packages/pyright-internal/src/languageServerBase.ts @@ -702,16 +702,13 @@ export abstract class LanguageServerBase implements LanguageServerInterface, Dis // not send config updates before this point. this._initialized = true; - if (!this.client.hasWorkspaceFoldersCapability) { - // If folder capability is not supported, initialize ones given by onInitialize. - this.updateSettingsForAllWorkspaces(); - return; + if (this.client.hasWorkspaceFoldersCapability) { + this._workspaceFoldersChangedDisposable = + this.connection.workspace.onDidChangeWorkspaceFolders(changeWorkspaceFolderHandler); } - this._workspaceFoldersChangedDisposable = - this.connection.workspace.onDidChangeWorkspaceFolders(changeWorkspaceFolderHandler); - this.dynamicFeatures.register(); + this.updateSettingsForAllWorkspaces(); } protected onDidChangeConfiguration(params: DidChangeConfigurationParams) { diff --git a/packages/pyright-internal/src/tests/languageServer.test.ts b/packages/pyright-internal/src/tests/languageServer.test.ts index 3ffcd092e554..ce6cd9bf3e5c 100644 --- a/packages/pyright-internal/src/tests/languageServer.test.ts +++ b/packages/pyright-internal/src/tests/languageServer.test.ts @@ -11,6 +11,7 @@ import { CancellationToken, CompletionRequest, ConfigurationItem, + DidChangeWorkspaceFoldersNotification, InitializedNotification, InitializeRequest, MarkupContent, @@ -26,7 +27,6 @@ import { DEFAULT_WORKSPACE_ROOT, getParseResults, hover, - initializeLanguageServer, openFile, PyrightServerInfo, runPyrightServer, @@ -68,45 +68,68 @@ describe(`Basic language server tests`, () => { await cleanupAfterAll(); }); - test('Basic Initialize', async () => { - const code = ` -// @filename: test.py -//// # empty file - `; - serverInfo = await runLanguageServer(DEFAULT_WORKSPACE_ROOT, code, /* callInitialize */ false); - - const initializeResult = await initializeLanguageServer(serverInfo); - - assert(initializeResult); - assert(initializeResult.capabilities.completionProvider?.resolveProvider); - }); - - test('Initialize without workspace folder support', async () => { + test.each([ + { name: 'capability disabled', capability: false, initFolders: 1, firstNotify: null, secondNotify: null }, + { name: '1 init, no notifications', capability: true, initFolders: 1, firstNotify: null, secondNotify: null }, + { name: '1 init, notify with 0', capability: true, initFolders: 1, firstNotify: 0, secondNotify: null }, + { name: '1 init, notify with 1', capability: true, initFolders: 1, firstNotify: 1, secondNotify: null }, + { name: '1 init, notify with 2', capability: true, initFolders: 1, firstNotify: 2, secondNotify: null }, + { name: '1 init, notify with 0 then 0', capability: true, initFolders: 1, firstNotify: 0, secondNotify: 0 }, + { name: '1 init, notify with 0 then 1', capability: true, initFolders: 1, firstNotify: 0, secondNotify: 1 }, + { name: '1 init, notify with 0 then 2', capability: true, initFolders: 1, firstNotify: 0, secondNotify: 2 }, + { name: '1 init, notify with 1 then 0', capability: true, initFolders: 1, firstNotify: 1, secondNotify: 0 }, + { name: '1 init, notify with 1 then 1', capability: true, initFolders: 1, firstNotify: 1, secondNotify: 1 }, + { name: '1 init, notify with 1 then 2', capability: true, initFolders: 1, firstNotify: 1, secondNotify: 2 }, + { name: '1 init, notify with 2 then 0', capability: true, initFolders: 1, firstNotify: 2, secondNotify: 0 }, + { name: '1 init, notify with 2 then 1', capability: true, initFolders: 1, firstNotify: 2, secondNotify: 1 }, + { name: '1 init, notify with 2 then 2', capability: true, initFolders: 1, firstNotify: 2, secondNotify: 2 }, + { name: '2 init, no notifications', capability: true, initFolders: 2, firstNotify: null, secondNotify: null }, + { name: '2 init, notify with 2', capability: true, initFolders: 2, firstNotify: 2, secondNotify: null }, + { name: '0 init, notify with 1', capability: true, initFolders: 0, firstNotify: 1, secondNotify: null }, + { name: '0 init, notify with 2', capability: true, initFolders: 0, firstNotify: 2, secondNotify: null }, + ])('workspace initialization: $name', async ({ capability, initFolders, firstNotify, secondNotify }) => { const code = ` // @filename: test.py //// import [|/*marker*/os|] `; - const info = await runLanguageServer(DEFAULT_WORKSPACE_ROOT, code, /* callInitialize */ false); - - // This will test clients with no folder and configuration support. + const info = await runLanguageServer(DEFAULT_WORKSPACE_ROOT, code, false); const params = info.getInitializeParams(); - params.capabilities.workspace!.workspaceFolders = false; - params.capabilities.workspace!.configuration = false; - - // Perform LSP Initialize/Initialized handshake. - const result = await info.connection.sendRequest(InitializeRequest.type, params, CancellationToken.None); - assert(result); + const folders = params.workspaceFolders!; + const folder2 = { name: 'workspace2', uri: 'file:///workspace2' }; + + params.capabilities.workspace!.workspaceFolders = capability; + if (initFolders === 0) { + params.workspaceFolders = []; + } else if (initFolders === 2) { + params.workspaceFolders = [...folders, folder2]; + } + await info.connection.sendRequest(InitializeRequest.type, params, CancellationToken.None); await info.connection.sendNotification(InitializedNotification.type, {}); - // Do simple hover request to verify our server works with a client that doesn't support - // workspace folder/configuration capabilities. + const getFoldersForNotify = (count: number) => { + if (count === 0) return []; + if (count === 1) return folders; + return [...folders, folder2]; + }; + + if (firstNotify !== null) { + await info.connection.sendNotification(DidChangeWorkspaceFoldersNotification.type, { + event: { added: getFoldersForNotify(firstNotify), removed: [] }, + }); + } + if (secondNotify !== null) { + await info.connection.sendNotification(DidChangeWorkspaceFoldersNotification.type, { + event: { added: getFoldersForNotify(secondNotify), removed: [] }, + }); + } + openFile(info, 'marker'); - const hoverResult = await hover(info, 'marker'); - assert(hoverResult); - assert(MarkupContent.is(hoverResult.contents)); - assert.strictEqual(hoverResult.contents.value, '```python\n(module) os\n```'); + const result = await hover(info, 'marker'); + assert(result && MarkupContent.is(result.contents)); + assert.strictEqual(result.contents.value, '```python\n(module) os\n```'); }); + test('Hover', async () => { const code = ` // @filename: test.py diff --git a/packages/pyright-internal/src/tests/lsp/languageServerTestUtils.ts b/packages/pyright-internal/src/tests/lsp/languageServerTestUtils.ts index 9f845db889bc..00416c90c24c 100644 --- a/packages/pyright-internal/src/tests/lsp/languageServerTestUtils.ts +++ b/packages/pyright-internal/src/tests/lsp/languageServerTestUtils.ts @@ -801,7 +801,7 @@ export async function openFile(info: PyrightServerInfo, markerName: string, text } export async function hover(info: PyrightServerInfo, markerName: string) { - const marker = info.testData.markerPositions.get('marker')!; + const marker = info.testData.markerPositions.get(markerName)!; const fileUri = marker.fileUri; const text = info.testData.files.find((d) => d.fileName === marker.fileName)!.content; const parseResult = getParseResults(text); diff --git a/packages/pyright-internal/src/tests/samples/conditionalExpr1.py b/packages/pyright-internal/src/tests/samples/conditionalExpr1.py new file mode 100644 index 000000000000..8c61ac760070 --- /dev/null +++ b/packages/pyright-internal/src/tests/samples/conditionalExpr1.py @@ -0,0 +1,106 @@ +# This sample tests type narrowing for conditional expressions (ternary operator) +# where the condition is narrowed such that one branch is known to be unreachable. + +from typing import assert_type + + +class Node: + pass + + +class Wrapper: + def __init__(self, child: Node): + self.child = child + + +class SpecialNode(Node): + pass + + +def func1(plan: Wrapper | None): + # After the guard, plan is known to be truthy (not None) and + # plan.child is known to be SpecialNode. + if not (plan and isinstance(plan.child, SpecialNode)): + return + + # This should be fine - direct assignment works. + ts1: SpecialNode = plan.child + + # This should also be fine - the else branch (None) is unreachable + # because plan is known to be truthy in this context. + ts2: SpecialNode = plan.child if plan else None + + # Also verify the inferred type. + assert_type(plan.child if plan else None, SpecialNode) + + +def func2(val: int | None): + # After this guard, val is known to be truthy (not None and not 0). + if not val: + return + + # The else branch is unreachable since val is known to be truthy. + # Using different types to make the test meaningful. + ts1: int = val if val else "fallback" + assert_type(val if val else "fallback", int) + + +def func3(val: str | None): + # After this guard, val is known to be None (falsy). + if val: + return + + # The if branch is unreachable since val is known to be falsy (None). + # Using different types to make the test meaningful - without pruning, + # this would be str | None instead of just None. + ts1: None = "unreachable" if val else None + assert_type("unreachable" if val else None, None) + + +def func4(val: int | None): + # After this guard, val is known to be not None (but could still be 0, which is falsy). + if val is None: + return + + # The else branch is still reachable since val could be 0 (falsy). + # This test verifies that we don't over-narrow. + ts1: int | str = val if val else "zero" + assert_type(val if val else "zero", int | str) + + +def func5(val: int | None): + # After this guard, val is known to be truthy (not None and not 0). + if not val: + return + + # The else branch is unreachable since val is known to be truthy. + # Using different types to make the test meaningful. + ts1: int = val if val else "fallback" + assert_type(val if val else "fallback", int) + + +def func6(val: list[int] | None): + # After this guard, val is known to be not None. + if val is None: + return + + # However, val could still be an empty list (falsy), so the else branch + # is still reachable. Both branches should contribute to the type. + ts1: list[int] | str = val if val else "empty" + assert_type(val if val else "empty", list[int] | str) + + +def func_bool_literal(): + # Test that the bool literal guard prevents over-narrowing for mutable variables. + maybe = True + # Both branches should remain since maybe could be reassigned. + ts: int | str = 1 if maybe else "no" + assert_type(1 if maybe else "no", int | str) + + +def func_bool_literal_not(): + # Test that the bool literal guard also applies to `not` expressions. + flag = True + # Both branches should remain since flag could be reassigned. + ts: int | str = 1 if not flag else "yes" + assert_type(1 if not flag else "yes", int | str) diff --git a/packages/pyright-internal/src/tests/typeEvaluator5.test.ts b/packages/pyright-internal/src/tests/typeEvaluator5.test.ts index bf35f9c0e94b..b1fe3aa6ad5d 100644 --- a/packages/pyright-internal/src/tests/typeEvaluator5.test.ts +++ b/packages/pyright-internal/src/tests/typeEvaluator5.test.ts @@ -315,6 +315,11 @@ test('Conditional1', () => { TestUtils.validateResults(analysisResults, 15); }); +test('ConditionalExpr1', () => { + const analysisResults = TestUtils.typeAnalyzeSampleFiles(['conditionalExpr1.py']); + TestUtils.validateResults(analysisResults, 0); +}); + test('TypePrinter1', () => { const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typePrinter1.py']); TestUtils.validateResults(analysisResults, 0); diff --git a/packages/pyright-internal/src/workspaceFactory.ts b/packages/pyright-internal/src/workspaceFactory.ts index 99bd0e7ea478..446b664b76c7 100644 --- a/packages/pyright-internal/src/workspaceFactory.ts +++ b/packages/pyright-internal/src/workspaceFactory.ts @@ -151,7 +151,11 @@ export class WorkspaceFactory implements IWorkspaceFactory { params.added.forEach((workspaceInfo) => { const uri = Uri.parse(workspaceInfo.uri, this._serviceProvider); - // Add the new workspace. + // Skip if workspace already exists (e.g., created during initialize) + if (this.getNonDefaultWorkspaces().some((w) => w.rootUri.equals(uri))) { + return; + } + this._add(uri, workspaceInfo.name, [WellKnownWorkspaceKinds.Regular]); });