From 0cab90b5ad88fe8b3b08a7f0ddc6b79ebc941662 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 6 Mar 2026 19:29:30 +0900 Subject: [PATCH 1/4] impl --- lsp/src/extension.ts | 40 ++++ lsp/src/overrideCodeLens.ts | 308 +++++++++++++++++++++++++++++ pyrefly/lib/state/lsp.rs | 74 +++++++ pyrefly/lib/test/lsp/definition.rs | 76 +++++++ 4 files changed, 498 insertions(+) create mode 100644 lsp/src/overrideCodeLens.ts diff --git a/lsp/src/extension.ts b/lsp/src/extension.ts index be948979ac..6b6397c41f 100644 --- a/lsp/src/extension.ts +++ b/lsp/src/extension.ts @@ -36,6 +36,11 @@ import {runDocstringFoldingCommand} from './docstring'; import {registerCodeLensCommands} from './code-lens'; import {registerHoverProvider} from './hover'; import {PythonEnvironment} from './python-environment'; +import { + OverrideCodeLensProvider, + navigateToOverrideTarget, + showOverrideTargets, +} from './overrideCodeLens'; import { triggerMsPythonRefreshLanguageServersIfInstalled, } from './extension-interop'; @@ -252,6 +257,7 @@ export async function activate( } pythonEnv = new PythonEnvironment(context); + const overrideCodeLensProvider = new OverrideCodeLensProvider(); launchSpec = await resolveExecutable( context.extensionUri, @@ -347,8 +353,21 @@ export async function activate( registerHoverProvider(context, () => client); } + context.subscriptions.push( + vscode.languages.registerCodeLensProvider( + [ + {scheme: 'file', language: 'python'}, + {scheme: 'untitled', language: 'python'}, + {scheme: 'vscode-notebook-cell', language: 'python'}, + {scheme: 'inmemory', language: 'python'}, + ], + overrideCodeLensProvider, + ), + ); + context.subscriptions.push( vscode.window.onDidChangeActiveTextEditor(async () => { + overrideCodeLensProvider.refresh(); await updateStatusBar(client); }), ); @@ -386,6 +405,7 @@ export async function activate( client.sendNotification(DidChangeConfigurationNotification.type, { settings: {}, }); + overrideCodeLensProvider.refresh(); } if ( event.affectsConfiguration('pyrefly.lspPath') || @@ -401,6 +421,14 @@ export async function activate( }), ); + context.subscriptions.push( + workspace.onDidChangeTextDocument(event => { + if (event.document.languageId === 'python') { + overrideCodeLensProvider.refresh(); + } + }), + ); + context.subscriptions.push( vscode.commands.registerCommand('pyrefly.restartClient', async () => { // Clear the output channel but don't dispose it @@ -510,6 +538,17 @@ export async function activate( ); registerCodeLensCommands(context, pythonEnv); + context.subscriptions.push( + vscode.commands.registerCommand( + 'pyrefly.navigateToOverrideTarget', + navigateToOverrideTarget, + ), + vscode.commands.registerCommand( + 'pyrefly.showOverrideTargets', + showOverrideTargets, + ), + ); + // When our extension is activated, make sure ms-python knows // TODO(kylei): remove this hack once ms-python has this behavior await triggerMsPythonRefreshLanguageServersIfInstalled(); @@ -537,6 +576,7 @@ export async function activate( } }); logServerVersion(); + overrideCodeLensProvider.refresh(); await updateStatusBar(client); const statusBarItem = getStatusBarItem(); diff --git a/lsp/src/overrideCodeLens.ts b/lsp/src/overrideCodeLens.ts new file mode 100644 index 0000000000..fbf4ad1f00 --- /dev/null +++ b/lsp/src/overrideCodeLens.ts @@ -0,0 +1,308 @@ +/** + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + * + * @format + */ + +import * as path from 'path'; +import * as vscode from 'vscode'; + +type MethodSymbol = { + range: vscode.Range; + position: vscode.Position; +}; + +type NavigationTarget = { + name: string; + uri: vscode.Uri; + range: vscode.Range; +}; + +function codeLensDisabled(document: vscode.TextDocument): boolean { + return vscode.workspace + .getConfiguration('python.pyrefly', document.uri) + .get>('disabledLanguageServices', {}) + .codeLens === true; +} + +function rangesEqual(left: vscode.Range, right: vscode.Range): boolean { + return left.start.isEqual(right.start) && left.end.isEqual(right.end); +} + +function locationsEqual( + left: vscode.Location, + rightUri: vscode.Uri, + rightRange: vscode.Range, +): boolean { + return left.uri.toString() === rightUri.toString() && rangesEqual(left.range, rightRange); +} + +function normalizeLocations( + result: vscode.Location | vscode.Location[] | vscode.LocationLink[] | undefined, +): vscode.Location[] { + if (!result) { + return []; + } + + const items = Array.isArray(result) ? result : [result]; + const locations = items.map(item => + 'targetUri' in item + ? new vscode.Location(item.targetUri, item.targetSelectionRange ?? item.targetRange) + : item, + ); + + const seen = new Set(); + return locations.filter(location => { + const key = `${location.uri.toString()}:${location.range.start.line}:${location.range.start.character}:${location.range.end.line}:${location.range.end.character}`; + if (seen.has(key)) { + return false; + } + seen.add(key); + return true; + }); +} + +function collectMethodSymbols( + symbols: readonly vscode.DocumentSymbol[], + enclosingClasses: readonly string[] = [], +): MethodSymbol[] { + const methods: MethodSymbol[] = []; + for (const symbol of symbols) { + if (symbol.kind === vscode.SymbolKind.Class) { + methods.push( + ...collectMethodSymbols(symbol.children, [...enclosingClasses, symbol.name]), + ); + continue; + } + if (symbol.kind !== vscode.SymbolKind.Method || enclosingClasses.length === 0) { + continue; + } + methods.push({ + range: symbol.selectionRange, + position: symbol.selectionRange.start, + }); + } + return methods; +} + +function findSymbolNameAtRange( + symbols: readonly vscode.DocumentSymbol[], + targetRange: vscode.Range, + enclosingClasses: readonly string[] = [], +): string | undefined { + for (const symbol of symbols) { + const nextEnclosingClasses = + symbol.kind === vscode.SymbolKind.Class + ? [...enclosingClasses, symbol.name] + : enclosingClasses; + if (rangesEqual(symbol.selectionRange, targetRange)) { + if (symbol.kind === vscode.SymbolKind.Method && nextEnclosingClasses.length > 0) { + return `${nextEnclosingClasses.join('.')}.${symbol.name}`; + } + return symbol.name; + } + const nested = findSymbolNameAtRange( + symbol.children, + targetRange, + nextEnclosingClasses, + ); + if (nested) { + return nested; + } + } + return undefined; +} + +async function getDocumentSymbols( + uri: vscode.Uri, + cache: Map>, +): Promise { + const key = uri.toString(); + let symbols = cache.get(key); + if (!symbols) { + symbols = Promise.resolve( + vscode.commands.executeCommand( + 'vscode.executeDocumentSymbolProvider', + uri, + ), + ).then(result => result ?? []); + cache.set(key, symbols); + } + return await symbols; +} + +async function describeLocation( + location: vscode.Location, + cache: Map>, +): Promise { + const symbols = await getDocumentSymbols(location.uri, cache); + const symbolName = findSymbolNameAtRange(symbols, location.range); + if (symbolName) { + return symbolName; + } + return `${path.basename(location.uri.fsPath)}:${location.range.start.line + 1}`; +} + +function singleTargetCommand(target: NavigationTarget): vscode.Command { + return { + title: '', + command: 'pyrefly.navigateToOverrideTarget', + arguments: [target], + }; +} + +function multiTargetCommand(targets: NavigationTarget[]): vscode.Command { + return { + title: '', + command: 'pyrefly.showOverrideTargets', + arguments: [targets], + }; +} + +function overrideLensTitle(targets: readonly NavigationTarget[]): string { + return targets.length === 1 + ? `$(arrow-up) Overrides ${targets[0].name}` + : `$(arrow-up) Overrides ${targets.length} parents`; +} + +function implementationLensTitle(targets: readonly NavigationTarget[]): string { + return targets.length === 1 + ? `$(arrow-down) Implemented by ${targets[0].name}` + : `$(arrow-down) ${targets.length} implementations`; +} + +export class OverrideCodeLensProvider implements vscode.CodeLensProvider { + private readonly onDidChangeEmitter = new vscode.EventEmitter(); + + public readonly onDidChangeCodeLenses = this.onDidChangeEmitter.event; + + public refresh(): void { + this.onDidChangeEmitter.fire(); + } + + public async provideCodeLenses( + document: vscode.TextDocument, + token: vscode.CancellationToken, + ): Promise { + if (document.languageId !== 'python' || codeLensDisabled(document)) { + return []; + } + + const rootSymbols = + (await vscode.commands.executeCommand( + 'vscode.executeDocumentSymbolProvider', + document.uri, + )) ?? []; + if (rootSymbols.length === 0) { + return []; + } + + const symbolCache = new Map>(); + symbolCache.set(document.uri.toString(), Promise.resolve(rootSymbols)); + + const methodSymbols = collectMethodSymbols(rootSymbols); + const codeLenses = await Promise.all( + methodSymbols.map(method => + this.getCodeLensesForMethod(document, method, symbolCache, token), + ), + ); + return codeLenses.flat(); + } + + private async getCodeLensesForMethod( + document: vscode.TextDocument, + method: MethodSymbol, + symbolCache: Map>, + token: vscode.CancellationToken, + ): Promise { + if (token.isCancellationRequested) { + return []; + } + + const [declarationResult, implementationResult] = await Promise.all([ + vscode.commands.executeCommand( + 'vscode.executeDeclarationProvider', + document.uri, + method.position, + ), + vscode.commands.executeCommand( + 'vscode.executeImplementationProvider', + document.uri, + method.position, + ), + ]); + + const declarationTargets = normalizeLocations(declarationResult).filter( + location => !locationsEqual(location, document.uri, method.range), + ); + const implementationTargets = normalizeLocations(implementationResult).filter( + location => !locationsEqual(location, document.uri, method.range), + ); + + const lenses: vscode.CodeLens[] = []; + if (declarationTargets.length > 0) { + const targets = await Promise.all( + declarationTargets.map(async location => ({ + name: await describeLocation(location, symbolCache), + uri: location.uri, + range: location.range, + })), + ); + const command = targets.length === 1 ? singleTargetCommand(targets[0]) : multiTargetCommand(targets); + command.title = overrideLensTitle(targets); + lenses.push(new vscode.CodeLens(method.range, command)); + } + + if (implementationTargets.length > 0) { + const targets = await Promise.all( + implementationTargets.map(async location => ({ + name: await describeLocation(location, symbolCache), + uri: location.uri, + range: location.range, + })), + ); + const command = + targets.length === 1 ? singleTargetCommand(targets[0]) : multiTargetCommand(targets); + command.title = implementationLensTitle(targets); + lenses.push(new vscode.CodeLens(method.range, command)); + } + + return lenses; + } +} + +export async function navigateToOverrideTarget( + target: NavigationTarget, +): Promise { + const document = await vscode.workspace.openTextDocument(target.uri); + const editor = await vscode.window.showTextDocument(document); + editor.selection = new vscode.Selection(target.range.start, target.range.start); + editor.revealRange(target.range, vscode.TextEditorRevealType.InCenter); +} + +export async function showOverrideTargets( + targets: NavigationTarget[], +): Promise { + if (targets.length === 0) { + return; + } + if (targets.length === 1) { + await navigateToOverrideTarget(targets[0]); + return; + } + + const pickedTarget = await vscode.window.showQuickPick( + targets.map(target => ({ + label: target.name, + description: path.basename(target.uri.fsPath), + target, + })), + {placeHolder: 'Select a parent or implementation to navigate to'}, + ); + if (pickedTarget) { + await navigateToOverrideTarget(pickedTarget.target); + } +} diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 8088d344cc..e2b9c0dbc6 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -3374,6 +3374,45 @@ impl<'a> Transaction<'a> { }) } + /// Finds overridden parent declarations for a method or class attribute definition. + /// + /// The answers index records override edges from a child declaration range to the + /// declaration ranges it overrides in the MRO. We use that mapping to make + /// go-to-declaration on an overriding member jump to the parent declaration instead + /// of back to the local implementation. + fn parent_declarations_for_definition( + &self, + handle: &Handle, + definition: &TextRangeWithModule, + ) -> Vec { + let Some(index) = self + .get_solutions(handle) + .and_then(|solutions| solutions.get_index()) + else { + return Vec::new(); + }; + let Some(parent_methods) = index + .lock() + .parent_methods_map + .get(&definition.range) + .cloned() + else { + return Vec::new(); + }; + + let mut seen = HashSet::new(); + let mut declarations = Vec::new(); + for (parent_module_path, parent_range) in parent_methods { + let Some(module) = module_info_for_path(self, &parent_module_path) else { + continue; + }; + if seen.insert((module.path().dupe(), parent_range)) { + declarations.push(TextRangeWithModule::new(module, parent_range)); + } + } + declarations + } + pub fn goto_declaration( &self, handle: &Handle, @@ -3391,6 +3430,23 @@ impl<'a> Transaction<'a> { }, )?; + let declarations = definitions + .iter() + .flat_map(|definition| { + self.parent_declarations_for_definition( + handle, + &TextRangeWithModule::new( + definition.module.dupe(), + definition.definition_range, + ), + ) + }) + .collect::>(); + + if !declarations.is_empty() { + return Ok(declarations); + } + Ok(definitions .into_vec() .into_map(|item| TextRangeWithModule::new(item.module, item.definition_range))) @@ -5360,6 +5416,24 @@ fn find_child_implementations_impl( child_implementations } +fn module_info_for_path(transaction: &Transaction<'_>, module_path: &ModulePath) -> Option { + let counterpart_path = match module_path.details() { + ModulePathDetails::Memory(path) => Some(ModulePath::filesystem((**path).clone())), + ModulePathDetails::FileSystem(path) => Some(ModulePath::memory(path.clone().to_path_buf())), + _ => None, + }; + + transaction.handles().into_iter().find_map(|handle| { + let path = handle.path(); + ((path == module_path) + || counterpart_path + .as_ref() + .is_some_and(|counterpart| path == counterpart)) + .then(|| transaction.get_module_info(&handle)) + .flatten() + }) +} + fn compute_transitive_rdeps_for_definition_impl( transaction: &mut T, sys_info: SysInfo, diff --git a/pyrefly/lib/test/lsp/definition.rs b/pyrefly/lib/test/lsp/definition.rs index 95e70b9732..2e67b86326 100644 --- a/pyrefly/lib/test/lsp/definition.rs +++ b/pyrefly/lib/test/lsp/definition.rs @@ -77,6 +77,30 @@ fn get_test_report_do_not_jump_through_renamed_import( } } +fn get_declaration_report(state: &State, handle: &Handle, position: TextSize) -> String { + let defs = state + .transaction() + .goto_declaration(handle, position) + .unwrap_or_default(); + if !defs.is_empty() { + defs.into_iter() + .map( + |TextRangeWithModule { + module: module_info, + range, + }| { + format!( + "Declaration Result:\n{}", + code_frame_of_source_at_range(module_info.contents(), range) + ) + }, + ) + .join("\n") + } else { + "Declaration Result: None".to_owned() + } +} + #[test] fn endpoint_definition_from_client_literal() { let code = r#" @@ -486,6 +510,58 @@ Definition Result: ); } +#[test] +fn overriding_method_declaration_jumps_to_parent_test() { + let code = r#" +class Base: + def f(self) -> int: + return 1 + +class Child(Base): + def f(self) -> int: +# ^ + return 2 +"#; + let report = get_batched_lsp_operations_report(&[("main", code)], get_declaration_report); + assert_eq!( + r#" +# main.py +7 | def f(self) -> int: + ^ +Declaration Result: +3 | def f(self) -> int: + ^ +"# + .trim(), + report.trim(), + ); +} + +#[test] +fn overriding_class_attribute_declaration_jumps_to_parent_test() { + let code = r#" +class Base: + value = 1 + +class Child(Base): + value = 2 +# ^ +"#; + let report = get_batched_lsp_operations_report(&[("main", code)], get_declaration_report); + assert_eq!( + r#" +# main.py +6 | value = 2 + ^ +Declaration Result: +3 | value = 1 + ^^^^^ +"# + .trim(), + report.trim(), + ); +} + #[test] fn narrow_test() { let code = r#" From 2562cbd49abe1691066de29d675959636f5c870a Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 7 Mar 2026 04:37:21 +0900 Subject: [PATCH 2/4] copilot --- lsp/src/extension.ts | 32 +++++++++++--- lsp/src/overrideCodeLens.ts | 61 ++++++++++++++++++++++---- pyrefly/lib/state/lsp.rs | 87 +++++++++++++++++++++++++------------ 3 files changed, 139 insertions(+), 41 deletions(-) diff --git a/lsp/src/extension.ts b/lsp/src/extension.ts index 6b6397c41f..7212fd60fe 100644 --- a/lsp/src/extension.ts +++ b/lsp/src/extension.ts @@ -258,6 +258,24 @@ export async function activate( pythonEnv = new PythonEnvironment(context); const overrideCodeLensProvider = new OverrideCodeLensProvider(); + let codeLensRefreshTimeout: NodeJS.Timeout | undefined; + const scheduleCodeLensRefresh = () => { + if (codeLensRefreshTimeout) { + clearTimeout(codeLensRefreshTimeout); + } + codeLensRefreshTimeout = setTimeout(() => { + codeLensRefreshTimeout = undefined; + overrideCodeLensProvider.refresh(); + }, 150); + }; + + context.subscriptions.push( + new vscode.Disposable(() => { + if (codeLensRefreshTimeout) { + clearTimeout(codeLensRefreshTimeout); + } + }), + ); launchSpec = await resolveExecutable( context.extensionUri, @@ -367,7 +385,7 @@ export async function activate( context.subscriptions.push( vscode.window.onDidChangeActiveTextEditor(async () => { - overrideCodeLensProvider.refresh(); + scheduleCodeLensRefresh(); await updateStatusBar(client); }), ); @@ -405,7 +423,7 @@ export async function activate( client.sendNotification(DidChangeConfigurationNotification.type, { settings: {}, }); - overrideCodeLensProvider.refresh(); + scheduleCodeLensRefresh(); } if ( event.affectsConfiguration('pyrefly.lspPath') || @@ -423,8 +441,11 @@ export async function activate( context.subscriptions.push( workspace.onDidChangeTextDocument(event => { - if (event.document.languageId === 'python') { - overrideCodeLensProvider.refresh(); + if ( + event.document.languageId === 'python' && + event.document === vscode.window.activeTextEditor?.document + ) { + scheduleCodeLensRefresh(); } }), ); @@ -463,6 +484,7 @@ export async function activate( ); } }); + scheduleCodeLensRefresh(); }), ); @@ -576,7 +598,7 @@ export async function activate( } }); logServerVersion(); - overrideCodeLensProvider.refresh(); + scheduleCodeLensRefresh(); await updateStatusBar(client); const statusBarItem = getStatusBarItem(); diff --git a/lsp/src/overrideCodeLens.ts b/lsp/src/overrideCodeLens.ts index fbf4ad1f00..6d61d4fe4d 100644 --- a/lsp/src/overrideCodeLens.ts +++ b/lsp/src/overrideCodeLens.ts @@ -10,6 +10,9 @@ import * as path from 'path'; import * as vscode from 'vscode'; +const MAX_METHOD_SYMBOLS = 200; +const CODE_LENS_CONCURRENCY = 12; + type MethodSymbol = { range: vscode.Range; position: vscode.Position; @@ -65,6 +68,12 @@ function normalizeLocations( }); } +function hasHierarchicalSymbols( + symbols: vscode.DocumentSymbol[] | vscode.SymbolInformation[] | undefined, +): symbols is vscode.DocumentSymbol[] { + return symbols !== undefined && (symbols.length === 0 || 'children' in symbols[0]); +} + function collectMethodSymbols( symbols: readonly vscode.DocumentSymbol[], enclosingClasses: readonly string[] = [], @@ -124,16 +133,42 @@ async function getDocumentSymbols( let symbols = cache.get(key); if (!symbols) { symbols = Promise.resolve( - vscode.commands.executeCommand( + vscode.commands.executeCommand< + vscode.DocumentSymbol[] | vscode.SymbolInformation[] + >( 'vscode.executeDocumentSymbolProvider', uri, ), - ).then(result => result ?? []); + ).then(result => (hasHierarchicalSymbols(result) ? result : [])); cache.set(key, symbols); } return await symbols; } +async function mapMethodsWithConcurrency( + methodSymbols: readonly MethodSymbol[], + token: vscode.CancellationToken, + mapper: (method: MethodSymbol) => Promise, +): Promise { + const results: T[] = new Array(methodSymbols.length); + let nextIndex = 0; + + const runNext = async (): Promise => { + while (!token.isCancellationRequested) { + const currentIndex = nextIndex; + nextIndex += 1; + if (currentIndex >= methodSymbols.length) { + return; + } + results[currentIndex] = await mapper(methodSymbols[currentIndex]); + } + }; + + const workerCount = Math.min(CODE_LENS_CONCURRENCY, methodSymbols.length); + await Promise.all(Array.from({length: workerCount}, () => runNext())); + return results; +} + async function describeLocation( location: vscode.Location, cache: Map>, @@ -191,11 +226,17 @@ export class OverrideCodeLensProvider implements vscode.CodeLensProvider { return []; } - const rootSymbols = - (await vscode.commands.executeCommand( + const rootSymbolsResult = + (await vscode.commands.executeCommand< + vscode.DocumentSymbol[] | vscode.SymbolInformation[] + >( 'vscode.executeDocumentSymbolProvider', document.uri, )) ?? []; + if (!hasHierarchicalSymbols(rootSymbolsResult) || rootSymbolsResult.length === 0) { + return []; + } + const rootSymbols = rootSymbolsResult; if (rootSymbols.length === 0) { return []; } @@ -204,10 +245,14 @@ export class OverrideCodeLensProvider implements vscode.CodeLensProvider { symbolCache.set(document.uri.toString(), Promise.resolve(rootSymbols)); const methodSymbols = collectMethodSymbols(rootSymbols); - const codeLenses = await Promise.all( - methodSymbols.map(method => - this.getCodeLensesForMethod(document, method, symbolCache, token), - ), + if (methodSymbols.length > MAX_METHOD_SYMBOLS) { + return []; + } + + const codeLenses = await mapMethodsWithConcurrency( + methodSymbols, + token, + method => this.getCodeLensesForMethod(document, method, symbolCache, token), ); return codeLenses.flat(); } diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index e2b9c0dbc6..774dc60bc8 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -3384,11 +3384,9 @@ impl<'a> Transaction<'a> { &self, handle: &Handle, definition: &TextRangeWithModule, + module_cache: &mut SmallMap>, ) -> Vec { - let Some(index) = self - .get_solutions(handle) - .and_then(|solutions| solutions.get_index()) - else { + let Some(index) = solutions_index_for_module(self, handle, &definition.module) else { return Vec::new(); }; let Some(parent_methods) = index @@ -3403,7 +3401,7 @@ impl<'a> Transaction<'a> { let mut seen = HashSet::new(); let mut declarations = Vec::new(); for (parent_module_path, parent_range) in parent_methods { - let Some(module) = module_info_for_path(self, &parent_module_path) else { + let Some(module) = module_info_for_path(self, module_cache, &parent_module_path) else { continue; }; if seen.insert((module.path().dupe(), parent_range)) { @@ -3429,27 +3427,26 @@ impl<'a> Transaction<'a> { ..Default::default() }, )?; - - let declarations = definitions - .iter() - .flat_map(|definition| { - self.parent_declarations_for_definition( - handle, - &TextRangeWithModule::new( - definition.module.dupe(), - definition.definition_range, - ), - ) - }) - .collect::>(); - - if !declarations.is_empty() { - return Ok(declarations); + let mut module_cache = SmallMap::new(); + let mut seen = HashSet::new(); + let mut results = Vec::new(); + for definition in definitions { + let base = + TextRangeWithModule::new(definition.module.dupe(), definition.definition_range); + let declarations = + self.parent_declarations_for_definition(handle, &base, &mut module_cache); + let declarations = if declarations.is_empty() { + vec![base] + } else { + declarations + }; + for declaration in declarations { + if seen.insert((declaration.module.path().dupe(), declaration.range)) { + results.push(declaration); + } + } } - - Ok(definitions - .into_vec() - .into_map(|item| TextRangeWithModule::new(item.module, item.definition_range))) + Ok(results) } /// Where a function-valued type was defined. @@ -5416,14 +5413,22 @@ fn find_child_implementations_impl( child_implementations } -fn module_info_for_path(transaction: &Transaction<'_>, module_path: &ModulePath) -> Option { +fn module_info_for_path( + transaction: &Transaction<'_>, + cache: &mut SmallMap>, + module_path: &ModulePath, +) -> Option { + if let Some(module) = cache.get(module_path) { + return module.dupe(); + } + let counterpart_path = match module_path.details() { ModulePathDetails::Memory(path) => Some(ModulePath::filesystem((**path).clone())), ModulePathDetails::FileSystem(path) => Some(ModulePath::memory(path.clone().to_path_buf())), _ => None, }; - transaction.handles().into_iter().find_map(|handle| { + let module = transaction.handles().into_iter().find_map(|handle| { let path = handle.path(); ((path == module_path) || counterpart_path @@ -5431,7 +5436,33 @@ fn module_info_for_path(transaction: &Transaction<'_>, module_path: &ModulePath) .is_some_and(|counterpart| path == counterpart)) .then(|| transaction.get_module_info(&handle)) .flatten() - }) + }); + cache.insert(module_path.dupe(), module.dupe()); + if let Some(counterpart_path) = counterpart_path { + cache.insert(counterpart_path, module.dupe()); + } + module +} + +fn solutions_index_for_module( + transaction: &Transaction<'_>, + handle: &Handle, + module: &Module, +) -> Option>> { + let counterpart_path = match module.path().details() { + ModulePathDetails::Memory(path) => Some(ModulePath::filesystem((**path).clone())), + ModulePathDetails::FileSystem(path) => Some(ModulePath::memory(path.clone().to_path_buf())), + _ => None, + }; + std::iter::once(module.path().dupe()) + .chain(counterpart_path) + .find_map(|path| { + self::Transaction::get_solutions( + transaction, + &Handle::new(module.name(), path, handle.sys_info().dupe()), + ) + .and_then(|solutions| solutions.get_index()) + }) } fn compute_transitive_rdeps_for_definition_impl( From 1957ad714ef03eb0e46792fbbf129abe99cf5a2f Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 30 Oct 2026 14:27:37 +0900 Subject: [PATCH 3/4] test --- lsp/src/test/override-code-lens.test.ts | 263 ++++++++++++++++++ pyrefly/lib/test/lsp/definition.rs | 94 ++++++- pyrefly/lib/test/lsp/implementation.rs | 61 ++++ .../lib/test/lsp/lsp_interaction/code_lens.rs | 233 ++++++++++++++++ 4 files changed, 647 insertions(+), 4 deletions(-) create mode 100644 lsp/src/test/override-code-lens.test.ts diff --git a/lsp/src/test/override-code-lens.test.ts b/lsp/src/test/override-code-lens.test.ts new file mode 100644 index 0000000000..dbc6d3b9a5 --- /dev/null +++ b/lsp/src/test/override-code-lens.test.ts @@ -0,0 +1,263 @@ +/** + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +import * as assert from 'assert'; +import * as vscode from 'vscode'; +import {LanguageClient, LSPErrorCodes, ResponseError} from 'vscode-languageclient/node'; +import {OverrideCodeLensProvider} from '../overrideCodeLens'; + +type NamedLocation = vscode.Location & {name?: string}; +type Targets = {parents: NamedLocation[]; implementations: NamedLocation[]}; +type Batch = {textDocument: {version: number}; ranges: vscode.Range[]; refreshSupport: boolean}; + +suite('OverrideCodeLensProvider', () => { + let symbols: vscode.DocumentSymbol[]; + let document: vscode.TextDocument; + let provider: OverrideCodeLensProvider; + let registrations: vscode.Disposable[]; + let requests: Batch[]; + let symbolQueries: number; + let notifications: Map void>; + let client: LanguageClient; + let respond: (batch: Batch, token: vscode.CancellationToken) => Promise; + + function symbol(name: string, kind: vscode.SymbolKind, line: number) { + const range = new vscode.Range(line, 0, line, name.length); + return new vscode.DocumentSymbol( + name, '', kind, + kind === vscode.SymbolKind.Method ? range : new vscode.Range(line, 0, 209, 80), + range, + ); + } + + setup(async () => { + symbols = []; + requests = []; + symbolQueries = 0; + notifications = new Map(); + respond = async batch => batch.ranges.map(() => ({parents: [], implementations: []})); + const scheme = 'override-code-lens-test'; + registrations = [ + vscode.workspace.registerTextDocumentContentProvider(scheme, { + provideTextDocumentContent: () => ' '.repeat(80).concat('\n').repeat(210), + }), + vscode.languages.registerDocumentSymbolProvider({scheme}, { + provideDocumentSymbols: () => { symbolQueries += 1; return symbols; }, + }), + ]; + document = await vscode.languages.setTextDocumentLanguage( + await vscode.workspace.openTextDocument(vscode.Uri.parse(`${scheme}:/methods.py`)), + 'python', + ); + client = { + onNotification: (method: string, handler: () => void) => { + notifications.set(method, handler); + return new vscode.Disposable(() => notifications.delete(method)); + }, + code2ProtocolConverter: { + asVersionedTextDocumentIdentifier: (doc: vscode.TextDocument) => ({ + uri: doc.uri.toString(), version: doc.version, + }), + asRange: (range: vscode.Range) => range, + }, + protocol2CodeConverter: {asLocation: (location: vscode.Location) => location}, + sendRequest: async (method: string, batch: Batch, token: vscode.CancellationToken) => { + assert.strictEqual(method, 'pyrefly/textDocument/overrideCodeLens'); + requests.push(batch); + return respond(batch, token); + }, + } as unknown as LanguageClient; + provider = new OverrideCodeLensProvider(() => client); + provider.connectClient(client); + }); + + teardown(() => { + provider.dispose(); + registrations.forEach(registration => registration.dispose()); + }); + + test('finds methods inside comment sections and class sections', async () => { + const method = symbol('f', vscode.SymbolKind.Method, 3); + const cls = symbol('Child', vscode.SymbolKind.Class, 1); + const section = symbol('Models', vscode.SymbolKind.String, 0); + const classSection = symbol('Methods', vscode.SymbolKind.String, 2); + classSection.children = [method]; + cls.children = [classSection]; + section.children = [cls]; + respond = async () => [{ + parents: [new vscode.Location(document.uri, new vscode.Range(5, 0, 5, 1))], + implementations: [], + }]; + const token = new vscode.CancellationTokenSource(); + try { + for (const root of [cls, section]) { + symbols = [root]; + const lenses = await provider.provideCodeLenses(document, token.token); + assert.strictEqual(lenses.length, 1); + assert.deepStrictEqual(lenses[0].range, method.selectionRange); + } + } finally { + token.dispose(); + } + }); + + test('sends one cancellable batch for 200 methods without overrides', async () => { + const cls = symbol('C', vscode.SymbolKind.Class, 0); + cls.children = Array.from({length: 200}, (_, i) => + symbol(`method_${i}`, vscode.SymbolKind.Method, i + 1), + ); + symbols = [cls]; + const token = new vscode.CancellationTokenSource(); + try { + assert.deepStrictEqual(await provider.provideCodeLenses(document, token.token), []); + assert.strictEqual(requests.length, 1); + assert.strictEqual(requests[0].ranges.length, 200); + assert.strictEqual(requests[0].textDocument.version, document.version); + assert.strictEqual(requests[0].refreshSupport, true); + } finally { + token.dispose(); + } + }); + + test('uses batch names without requesting symbols for 500 targets', async () => { + const cls = symbol('Base', vscode.SymbolKind.Class, 0); + cls.children = [symbol('f', vscode.SymbolKind.Method, 1)]; + symbols = [cls]; + respond = async () => [{parents: [], implementations: Array.from({length: 500}, (_, i) => + Object.assign(new vscode.Location( + vscode.Uri.parse(`override-code-lens-test:/child_${i}.py`), new vscode.Range(0, 0, 0, 1), + ), {name: `Child${i}.f`}), + )}]; + const token = new vscode.CancellationTokenSource(); + try { + const lenses = await provider.provideCodeLenses(document, token.token); + assert.strictEqual(lenses.length, 1); + assert.strictEqual(lenses[0].command?.title, '$(arrow-down) 500 implementations'); + assert.strictEqual(lenses[0].command?.arguments?.[0][499].name, 'Child499.f'); + assert.strictEqual(requests.length, 1); + assert.strictEqual(symbolQueries, 1); + token.cancel(); + } finally { + token.dispose(); + } + }); + + test('retains lenses and retries server cancellation', async () => { + const cls = symbol('Child', vscode.SymbolKind.Class, 0); + cls.children = [symbol('f', vscode.SymbolKind.Method, 1)]; + symbols = [cls]; + const target = Object.assign( + new vscode.Location(document.uri, new vscode.Range(5, 0, 5, 1)), {name: 'Base.f'}, + ); + const success = async () => [{parents: [target], implementations: []}]; + const token = new vscode.CancellationTokenSource(); + try { + respond = success; + const lenses = await provider.provideCodeLenses(document, token.token); + for (const code of [LSPErrorCodes.RequestCancelled, LSPErrorCodes.ServerCancelled, + LSPErrorCodes.ContentModified]) { + const refresh = new Promise(resolve => { + const subscription = provider.onDidChangeCodeLenses(() => { + subscription.dispose(); resolve(); + }); + }); + respond = async () => { throw new ResponseError(code, 'Analysis changed'); }; + assert.deepStrictEqual(await provider.provideCodeLenses(document, token.token), lenses); + await refresh; + respond = success; + assert.deepStrictEqual(await provider.provideCodeLenses(document, token.token), lenses); + } + } finally { + token.dispose(); + } + }); + + test('limits retries when the server continues to cancel', async () => { + const cls = symbol('C', vscode.SymbolKind.Class, 0); + cls.children = [symbol('f', vscode.SymbolKind.Method, 1)]; + symbols = [cls]; + respond = async () => { throw new ResponseError(LSPErrorCodes.ServerCancelled, 'Rechecking'); }; + let refreshCount = 0; + const subscription = provider.onDidChangeCodeLenses(() => { refreshCount += 1; }); + const token = new vscode.CancellationTokenSource(); + try { + for (let i = 0; i < 3; i += 1) { + const refreshed = new Promise(resolve => { + const event = provider.onDidChangeCodeLenses(() => { event.dispose(); resolve(); }); + }); + assert.deepStrictEqual(await provider.provideCodeLenses(document, token.token), []); + await refreshed; + } + assert.deepStrictEqual(await provider.provideCodeLenses(document, token.token), []); + await new Promise(resolve => setTimeout(resolve, 200)); + assert.strictEqual(requests.length, 4); + assert.strictEqual(refreshCount, 3); + } finally { + subscription.dispose(); + token.dispose(); + } + }); + + test('refreshes on server invalidation and reconnects the notification handler', async () => { + const cls = symbol('C', vscode.SymbolKind.Class, 0); + cls.children = [symbol('f', vscode.SymbolKind.Method, 1)]; + symbols = [cls]; + let started!: () => void; + const entered = new Promise(resolve => { started = resolve; }); + let requestToken: vscode.CancellationToken | undefined; + respond = async (_, token) => { + requestToken = token; + started(); + return new Promise(resolve => token.onCancellationRequested(() => resolve([]))); + }; + provider.connectClient(client); + assert.strictEqual(notifications.size, 1); + const token = new vscode.CancellationTokenSource(); + const pending = provider.provideCodeLenses(document, token.token); + await entered; + const refreshed = new Promise(resolve => { + const subscription = provider.onDidChangeCodeLenses(() => { + subscription.dispose(); resolve(); + }); + }); + notifications.get('pyrefly/overrideCodeLens/refresh')!(); + assert.strictEqual(requestToken?.isCancellationRequested, true); + assert.deepStrictEqual(await pending, []); + await refreshed; + token.dispose(); + }); + + test('propagates cancellation and cancels superseded batches', async () => { + const cls = symbol('C', vscode.SymbolKind.Class, 0); + cls.children = [symbol('f', vscode.SymbolKind.Method, 1)]; + symbols = [cls]; + for (const cancel of ['token', 'refresh', 'supersede'] as const) { + let started!: () => void; + const entered = new Promise(resolve => { started = resolve; }); + let requestToken: vscode.CancellationToken | undefined; + respond = async (_, token) => { + requestToken = token; + started(); + return new Promise(resolve => token.onCancellationRequested(() => resolve([]))); + }; + const token = new vscode.CancellationTokenSource(); + const pending = provider.provideCodeLenses(document, token.token); + await entered; + if (cancel === 'token') { + token.cancel(); + } else if (cancel === 'refresh') { + provider.refresh(); + } else { + respond = async () => []; + await provider.provideCodeLenses(document, token.token); + } + assert.strictEqual(requestToken?.isCancellationRequested, true); + assert.deepStrictEqual(await pending, []); + token.dispose(); + } + }); +}); diff --git a/pyrefly/lib/test/lsp/definition.rs b/pyrefly/lib/test/lsp/definition.rs index 2e67b86326..608ede55b3 100644 --- a/pyrefly/lib/test/lsp/definition.rs +++ b/pyrefly/lib/test/lsp/definition.rs @@ -7,6 +7,7 @@ use std::fs; use std::slice; +use std::sync::Arc; use itertools::Itertools as _; use pretty_assertions::assert_eq; @@ -23,6 +24,7 @@ use tempfile::TempDir; use crate::config::config::ConfigFile; use crate::config::finder::ConfigFinder; +use crate::state::load::FileContents; use crate::state::lsp::FindPreference; use crate::state::require::Require; use crate::state::state::State; @@ -511,7 +513,7 @@ Definition Result: } #[test] -fn overriding_method_declaration_jumps_to_parent_test() { +fn overriding_method_declaration_stays_at_definition_test() { let code = r#" class Base: def f(self) -> int: @@ -529,7 +531,7 @@ class Child(Base): 7 | def f(self) -> int: ^ Declaration Result: -3 | def f(self) -> int: +7 | def f(self) -> int: ^ "# .trim(), @@ -538,7 +540,7 @@ Declaration Result: } #[test] -fn overriding_class_attribute_declaration_jumps_to_parent_test() { +fn overriding_class_attribute_declaration_stays_at_definition_test() { let code = r#" class Base: value = 1 @@ -554,7 +556,7 @@ class Child(Base): 6 | value = 2 ^ Declaration Result: -3 | value = 1 +6 | value = 2 ^^^^^ "# .trim(), @@ -562,6 +564,90 @@ Declaration Result: ); } +#[test] +fn parent_declarations_include_special_methods_and_libraries() { + let code = r#" +from collections.abc import Sized + +class Base: + def __init__(self): pass + +class Child(Base, Sized): + def __init__(self): pass + def __len__(self) -> int: return 0 +"#; + let mut env = TestEnv::new(); + env.add("main", code); + let (state, handle) = env.to_state(); + let transaction = state.transaction(); + for name in ["__init__", "__len__"] { + let start = code.rfind(name).unwrap() as u32; + let range = TextRange::new( + TextSize::new(start), + TextSize::new(start + name.len() as u32), + ); + let parents = transaction.parent_declarations(&handle("main"), range); + assert!(!parents.is_empty(), "Missing parent for {name}"); + assert!( + parents + .iter() + .all(|parent| parent.module.code_at(parent.range) == name) + ); + if name == "__init__" { + assert_eq!( + parents[0].range.start(), + TextSize::new(code.find(name).unwrap() as u32) + ); + } else { + assert_eq!(parents[0].module.name().as_str(), "typing"); + } + } +} + +#[test] +fn parent_declarations_keep_the_declaring_snapshot() { + let root = TempDir::new().unwrap(); + let path = root.path().join("base.py"); + let disk = "class Base:\r\n def f(self): pass\r\n"; + let memory = disk.replace("\r\n", "\n"); + fs::write(&path, disk).unwrap(); + let code = "from base import Base\nclass Child(Base):\n def f(self): pass\n"; + let mut env = TestEnv::new(); + env.add_real_path("base", path.clone()); + env.add("main", code); + let (state, handle) = env.to_state(); + let disk_handle = handle("base"); + let memory_handle = Handle::new( + disk_handle.module(), + ModulePath::memory(path.clone()), + *disk_handle.sys_info(), + ); + let mut transaction = state.new_committable_transaction(Require::Everything, None); + transaction.as_mut().set_memory(vec![( + path, + Some(Arc::new(FileContents::from_source(memory))), + )]); + transaction + .as_mut() + .run(&[memory_handle], Require::Everything, None); + state.commit_transaction(transaction, None); + let start = code.rfind("f(self)").unwrap() as u32; + let range = TextRange::new(TextSize::new(start), TextSize::new(start + 1)); + let parents = state + .transaction() + .parent_declarations(&handle("main"), range); + assert_eq!(parents.len(), 1); + assert_eq!(parents[0].module.code_at(parents[0].range), "f"); + assert_eq!( + parents[0] + .module + .to_lsp_range(parents[0].range) + .start + .character, + 8 + ); +} + #[test] fn narrow_test() { let code = r#" diff --git a/pyrefly/lib/test/lsp/implementation.rs b/pyrefly/lib/test/lsp/implementation.rs index 3768ad53b0..b25cfe9644 100644 --- a/pyrefly/lib/test/lsp/implementation.rs +++ b/pyrefly/lib/test/lsp/implementation.rs @@ -9,11 +9,13 @@ use itertools::Itertools as _; use pretty_assertions::assert_eq; use pyrefly_build::handle::Handle; use pyrefly_python::module::TextRangeWithModule; +use ruff_text_size::TextRange; use ruff_text_size::TextSize; use vec1::Vec1; use crate::state::lsp::FindPreference; use crate::state::state::State; +use crate::test::util::TestEnv; use crate::test::util::code_frame_of_source_at_range; use crate::test::util::get_batched_lsp_operations_report; @@ -139,3 +141,62 @@ Implementation Result: report.trim(), ); } + +#[test] +fn override_code_lens_batches_methods_across_modules() { + let base = "class Base:\n def f(self): pass\n def g(self): pass\n"; + let child = "from base import Base\nclass Child(Base):\n def f(self): pass\n"; + let mut env = TestEnv::new(); + env.add("base", base); + env.add("child", child); + let (state, handle) = env.to_state(); + let ranges = ["f(self)", "g(self)"].map(|name| { + let start = TextSize::new(base.find(name).unwrap() as u32); + TextRange::new(start, start + TextSize::new(1)) + }); + let targets = state + .cancellable_transaction() + .override_code_lens_targets(&handle("base"), &ranges) + .unwrap(); + assert_eq!(targets.len(), 2); + assert!(targets.iter().all(|target| target.parents.is_empty())); + assert_eq!(targets[0].implementations.len(), 1); + let implementation = &targets[0].implementations[0]; + assert_eq!(implementation.module.name().as_str(), "child"); + assert_eq!(implementation.module.code_at(implementation.range), "f"); + assert!(targets[1].implementations.is_empty()); +} + +#[test] +fn override_code_lens_batches_methods_without_overrides() { + let mut code = "class C:\n".to_owned(); + for i in 0..200 { + code.push_str(&format!(" def method_{i}(self): pass\n")); + } + let ranges = (0..200) + .map(|i| { + let name = format!("method_{i}"); + let start = TextSize::new(code.find(&format!("{name}(self)")).unwrap() as u32); + TextRange::new(start, start + TextSize::new(name.len() as u32)) + }) + .collect::>(); + let mut env = TestEnv::new(); + env.add("main", &code); + let (state, handle) = env.to_state(); + let mut transaction = state.cancellable_transaction(); + let targets = transaction + .override_code_lens_targets(&handle("main"), &ranges) + .unwrap(); + assert_eq!(targets.len(), 200); + assert!( + targets + .iter() + .all(|target| target.parents.is_empty() && target.implementations.is_empty()) + ); + transaction.get_cancellation_handle().cancel(); + assert!( + transaction + .override_code_lens_targets(&handle("main"), &ranges) + .is_err() + ); +} diff --git a/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs b/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs index f177a31669..895d74fe24 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs @@ -5,14 +5,23 @@ * LICENSE file in the root directory of this source tree. */ +use std::fs; + use lsp_types::CodeLens; use lsp_types::CodeLensRequest; use lsp_types::Uri; +use pyrefly_lsp_test::IndexingMode; +use pyrefly_lsp_test::LspArgs; +use pyrefly_lsp_test::Message; use pyrefly_lsp_test::object_model::InitializeSettings; use pyrefly_lsp_test::object_model::LspInteraction; +use pyrefly_lsp_test::object_model::LspInteractionArgs; use serde_json::Value; use serde_json::json; +use tempfile::TempDir; +use crate::lsp::non_wasm::code_lens::OverrideCodeLensRequest; +use crate::lsp::non_wasm::code_lens::OverrideCodeLensTargets; use crate::test::lsp::lsp_interaction::util::get_test_files_root; fn runnable_code_lens_config() -> serde_json::Value { @@ -206,3 +215,227 @@ fn test_code_lens_disabled_by_default() { interaction.shutdown().unwrap(); } + +#[test] +fn test_override_code_lens_batch_and_document_version() { + let root = get_test_files_root(); + let mut interaction = LspInteraction::new(); + let test_root = root.path().join("code_lens"); + interaction.set_root(test_root.clone()); + interaction.initialize(Default::default()).unwrap(); + let uri = Uri::from_file_path(test_root.join("overrides.py")).unwrap(); + interaction.client.did_open_uri( + &uri, + "python", + "class Base:\n def f(self): pass\nclass Child(Base):\n def f(self): pass\n", + ); + let params = json!({ + "textDocument": {"uri": uri.to_string(), "version": 1}, + "ranges": [ + {"start": {"line": 1, "character": 8}, "end": {"line": 1, "character": 9}}, + {"start": {"line": 3, "character": 8}, "end": {"line": 3, "character": 9}} + ] + }); + interaction + .client + .send_request::(params.clone()) + .expect_response_with(|targets: Vec| { + targets.len() == 2 + && targets[0].parents.is_empty() + && targets[0].implementations.len() == 1 + && targets[0].implementations[0].location.range.start.line == 3 + && targets[0].implementations[0].name == "Child.f" + && targets[1].parents.len() == 1 + && targets[1].parents[0].location.range.start.line == 1 + && targets[1].parents[0].name == "Base.f" + && targets[1].implementations.is_empty() + }) + .unwrap(); + let mut stale_params = params; + stale_params["textDocument"]["version"] = json!(0); + interaction + .client + .send_request::(stale_params) + .expect_response_error_code(-32800) + .unwrap(); + interaction.shutdown().unwrap(); +} + +#[test] +fn test_override_code_lens_notebook_cell_versions() { + let root = get_test_files_root(); + let mut interaction = LspInteraction::new(); + interaction.set_root(root.path().to_path_buf()); + interaction.initialize(Default::default()).unwrap(); + interaction.open_notebook( + "overrides.ipynb", + vec!["class Base:\n def f(self): pass\nclass Child(Base):\n def f(self): pass\n"], + ); + let uri = interaction.cell_uri("overrides.ipynb", "cell1"); + let mut params = json!({ + "textDocument": {"uri": uri.to_string(), "version": 1}, + "ranges": [{"start": {"line": 3, "character": 8}, "end": {"line": 3, "character": 9}}] + }); + for notebook_version in [1, 2] { + if notebook_version == 2 { + interaction.change_notebook( + "overrides.ipynb", + 2, + json!({"metadata": {"language_info": {"name": "python"}}}), + ); + } + interaction + .client + .send_request::(params.clone()) + .expect_response_with(|targets: Vec| { + targets.len() == 1 + && targets[0].parents.len() == 1 + && targets[0].parents[0].name == "Base.f" + && targets[0].parents[0].location.uri == uri + && targets[0].parents[0].location.range.start.line == 1 + }) + .unwrap(); + } + // The cell text version changes independently of the notebook version. + interaction.change_notebook("overrides.ipynb", 3, json!({"cells": {"textContent": [{ + "document": {"uri": uri.to_string(), "version": 7}, + "changes": [{"text": "class Base:\n def f(self): pass\nclass Child(Base):\n def f(self): pass\n"}] + }]}})); + interaction + .client + .send_request::(params.clone()) + .expect_response_error_code(-32800) + .unwrap(); + params["textDocument"]["version"] = json!(7); + interaction + .client + .send_request::(params) + .expect_response_with(|targets: Vec| { + targets.len() == 1 && targets[0].parents.len() == 1 + }) + .unwrap(); + interaction.shutdown().unwrap(); +} + +#[test] +fn test_override_code_lens_refresh_after_subclass_change() { + let root = get_test_files_root(); + let mut interaction = LspInteraction::new(); + let test_root = root.path().join("code_lens"); + interaction.set_root(test_root.clone()); + interaction.initialize(Default::default()).unwrap(); + let uri = Uri::from_file_path(test_root.join("base.py")).unwrap(); + interaction + .client + .did_open_uri(&uri, "python", "class Base:\n def f(self): pass\n"); + interaction + .client + .send_request::(json!({ + "textDocument": {"uri": uri.to_string(), "version": 1}, + "refreshSupport": true, + "ranges": [{"start": {"line": 1, "character": 8}, "end": {"line": 1, "character": 9}}] + })) + .expect_response_with(|_: Vec| true) + .unwrap(); + let child_uri = Uri::from_file_path(test_root.join("child.py")).unwrap(); + interaction.client.did_open_uri( + &child_uri, + "python", + "from base import Base\nclass Child(Base):\n def f(self): pass\n", + ); + interaction.client.did_change( + "child.py", + "from base import Base\nclass Child(Base):\n def f(self): return 1\n", + ); + interaction + .client + .expect_message("override CodeLens invalidation", |message| match message { + Message::Notification(notification) + if notification.method == "pyrefly/overrideCodeLens/refresh" => + { + Some(Ok(())) + } + _ => None, + }) + .unwrap(); + interaction.shutdown().unwrap(); +} + +#[test] +fn test_override_code_lens_refresh_after_indexed_file_change() { + let root = TempDir::new().unwrap(); + fs::write( + root.path().join("pyrefly.toml"), + "project-includes = ['*.py']\n", + ) + .unwrap(); + fs::write( + root.path().join("base.py"), + "class Base:\n def f(self): pass\n", + ) + .unwrap(); + let child = root.path().join("child.py"); + fs::write( + &child, + "from base import Base\nclass Child(Base):\n def f(self): pass\n", + ) + .unwrap(); + let uri = Uri::from_file_path(root.path().join("base.py")).unwrap(); + let scope = Uri::from_file_path(root.path()).unwrap(); + let mut interaction = LspInteraction::new_with_args(LspInteractionArgs { + args: LspArgs { + indexing_mode: IndexingMode::LazyBlocking, + ..LspInteractionArgs::default().args + }, + ..Default::default() + }); + interaction.set_root(root.path().to_path_buf()); + interaction + .initialize(InitializeSettings { + workspace_folders: Some(vec![("test".to_owned(), scope)]), + configuration: Some(None), + ..Default::default() + }) + .unwrap(); + interaction.client.did_open("base.py"); + let params = json!({ + "textDocument": {"uri": uri.to_string(), "version": 1}, + "refreshSupport": true, + "ranges": [{"start": {"line": 1, "character": 8}, "end": {"line": 1, "character": 9}}] + }); + interaction + .client + .send_request::(params.clone()) + .expect_response_with(|targets: Vec| { + targets.len() == 1 && targets[0].implementations.len() == 1 + }) + .unwrap(); + fs::write( + &child, + "from base import Base\nclass Child(Base):\n def g(self): pass\n", + ) + .unwrap(); + interaction.client.file_modified("child.py"); + interaction + .client + .expect_message( + "CodeLens invalidation after indexing", + |message| match message { + Message::Notification(notification) + if notification.method == "pyrefly/overrideCodeLens/refresh" => + { + Some(Ok(())) + } + _ => None, + }, + ) + .unwrap(); + interaction + .client + .send_request::(params) + .expect_response_with(|targets: Vec| { + targets.len() == 1 && targets[0].implementations.is_empty() + }) + .unwrap(); + interaction.shutdown().unwrap(); +} From 91d697791065c3d171a37f6c20fb4a120a33938f Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 30 Oct 2026 14:27:51 +0900 Subject: [PATCH 4/4] fix --- crates/pyrefly_util/src/task_heap.rs | 2 +- lsp/src/extension.ts | 30 +- lsp/src/overrideCodeLens.ts | 391 +++++++++++--------------- pyrefly/lib/lsp/non_wasm/code_lens.rs | 50 ++++ pyrefly/lib/lsp/non_wasm/server.rs | 241 +++++++++++++++- pyrefly/lib/state/lsp.rs | 232 ++++++++------- pyrefly/lib/state/notebook.rs | 16 +- 7 files changed, 607 insertions(+), 355 deletions(-) diff --git a/crates/pyrefly_util/src/task_heap.rs b/crates/pyrefly_util/src/task_heap.rs index ed3e8d6867..6c9546ddf4 100644 --- a/crates/pyrefly_util/src/task_heap.rs +++ b/crates/pyrefly_util/src/task_heap.rs @@ -23,7 +23,7 @@ use crate::lock::Mutex; #[derive(Debug)] pub struct Cancelled; -#[derive(Clone, Dupe)] +#[derive(Clone, Default, Dupe)] pub struct CancellationHandle(Arc); impl CancellationHandle { diff --git a/lsp/src/extension.ts b/lsp/src/extension.ts index 7212fd60fe..dd05604c71 100644 --- a/lsp/src/extension.ts +++ b/lsp/src/extension.ts @@ -47,6 +47,7 @@ import { import {describeError, resolveExecutable} from './lsp-path'; let client: LanguageClient; +let overrideCodeLensProvider: OverrideCodeLensProvider; let outputChannel: vscode.OutputChannel; let traceOutputChannel: vscode.OutputChannel; let inferOutputChannel: vscode.OutputChannel; @@ -101,6 +102,7 @@ async function replaceClient(): Promise { launchSpec, clientOptions, ); + overrideCodeLensProvider.connectClient(client); try { await client.start(); logServerVersion(); @@ -257,25 +259,9 @@ export async function activate( } pythonEnv = new PythonEnvironment(context); - const overrideCodeLensProvider = new OverrideCodeLensProvider(); - let codeLensRefreshTimeout: NodeJS.Timeout | undefined; - const scheduleCodeLensRefresh = () => { - if (codeLensRefreshTimeout) { - clearTimeout(codeLensRefreshTimeout); - } - codeLensRefreshTimeout = setTimeout(() => { - codeLensRefreshTimeout = undefined; - overrideCodeLensProvider.refresh(); - }, 150); - }; - - context.subscriptions.push( - new vscode.Disposable(() => { - if (codeLensRefreshTimeout) { - clearTimeout(codeLensRefreshTimeout); - } - }), - ); + overrideCodeLensProvider = new OverrideCodeLensProvider(() => client); + context.subscriptions.push(overrideCodeLensProvider); + const scheduleCodeLensRefresh = () => overrideCodeLensProvider.refresh(); launchSpec = await resolveExecutable( context.extensionUri, @@ -367,6 +353,7 @@ export async function activate( serverOptions, clientOptions, ); + overrideCodeLensProvider.connectClient(client); if (supportsHoverVerbosity) { registerHoverProvider(context, () => client); } @@ -441,10 +428,7 @@ export async function activate( context.subscriptions.push( workspace.onDidChangeTextDocument(event => { - if ( - event.document.languageId === 'python' && - event.document === vscode.window.activeTextEditor?.document - ) { + if (event.document.languageId === 'python') { scheduleCodeLensRefresh(); } }), diff --git a/lsp/src/overrideCodeLens.ts b/lsp/src/overrideCodeLens.ts index 6d61d4fe4d..4cdb17ef30 100644 --- a/lsp/src/overrideCodeLens.ts +++ b/lsp/src/overrideCodeLens.ts @@ -9,14 +9,14 @@ import * as path from 'path'; import * as vscode from 'vscode'; +import {LanguageClient, Location, LSPErrorCodes, ResponseError} from 'vscode-languageclient/node'; const MAX_METHOD_SYMBOLS = 200; -const CODE_LENS_CONCURRENCY = 12; +const REFRESH_DELAY_MS = 150; +const MAX_SERVER_RETRIES = 3; +const REFRESH_METHOD = 'pyrefly/overrideCodeLens/refresh'; -type MethodSymbol = { - range: vscode.Range; - position: vscode.Position; -}; +type NamedLocation = Location & {name?: string}; type NavigationTarget = { name: string; @@ -31,156 +31,32 @@ function codeLensDisabled(document: vscode.TextDocument): boolean { .codeLens === true; } -function rangesEqual(left: vscode.Range, right: vscode.Range): boolean { - return left.start.isEqual(right.start) && left.end.isEqual(right.end); -} - -function locationsEqual( - left: vscode.Location, - rightUri: vscode.Uri, - rightRange: vscode.Range, -): boolean { - return left.uri.toString() === rightUri.toString() && rangesEqual(left.range, rightRange); -} - -function normalizeLocations( - result: vscode.Location | vscode.Location[] | vscode.LocationLink[] | undefined, -): vscode.Location[] { - if (!result) { - return []; - } - - const items = Array.isArray(result) ? result : [result]; - const locations = items.map(item => - 'targetUri' in item - ? new vscode.Location(item.targetUri, item.targetSelectionRange ?? item.targetRange) - : item, - ); - - const seen = new Set(); - return locations.filter(location => { - const key = `${location.uri.toString()}:${location.range.start.line}:${location.range.start.character}:${location.range.end.line}:${location.range.end.character}`; - if (seen.has(key)) { - return false; - } - seen.add(key); - return true; - }); -} - function hasHierarchicalSymbols( symbols: vscode.DocumentSymbol[] | vscode.SymbolInformation[] | undefined, ): symbols is vscode.DocumentSymbol[] { return symbols !== undefined && (symbols.length === 0 || 'children' in symbols[0]); } -function collectMethodSymbols( +function collectMethodRanges( symbols: readonly vscode.DocumentSymbol[], - enclosingClasses: readonly string[] = [], -): MethodSymbol[] { - const methods: MethodSymbol[] = []; + inClass = false, +): vscode.Range[] { + const methods: vscode.Range[] = []; for (const symbol of symbols) { - if (symbol.kind === vscode.SymbolKind.Class) { + if (symbol.kind === vscode.SymbolKind.Method && inClass) { + methods.push(symbol.selectionRange); + } else { methods.push( - ...collectMethodSymbols(symbol.children, [...enclosingClasses, symbol.name]), + ...collectMethodRanges( + symbol.children, + inClass || symbol.kind === vscode.SymbolKind.Class, + ), ); - continue; } - if (symbol.kind !== vscode.SymbolKind.Method || enclosingClasses.length === 0) { - continue; - } - methods.push({ - range: symbol.selectionRange, - position: symbol.selectionRange.start, - }); } return methods; } -function findSymbolNameAtRange( - symbols: readonly vscode.DocumentSymbol[], - targetRange: vscode.Range, - enclosingClasses: readonly string[] = [], -): string | undefined { - for (const symbol of symbols) { - const nextEnclosingClasses = - symbol.kind === vscode.SymbolKind.Class - ? [...enclosingClasses, symbol.name] - : enclosingClasses; - if (rangesEqual(symbol.selectionRange, targetRange)) { - if (symbol.kind === vscode.SymbolKind.Method && nextEnclosingClasses.length > 0) { - return `${nextEnclosingClasses.join('.')}.${symbol.name}`; - } - return symbol.name; - } - const nested = findSymbolNameAtRange( - symbol.children, - targetRange, - nextEnclosingClasses, - ); - if (nested) { - return nested; - } - } - return undefined; -} - -async function getDocumentSymbols( - uri: vscode.Uri, - cache: Map>, -): Promise { - const key = uri.toString(); - let symbols = cache.get(key); - if (!symbols) { - symbols = Promise.resolve( - vscode.commands.executeCommand< - vscode.DocumentSymbol[] | vscode.SymbolInformation[] - >( - 'vscode.executeDocumentSymbolProvider', - uri, - ), - ).then(result => (hasHierarchicalSymbols(result) ? result : [])); - cache.set(key, symbols); - } - return await symbols; -} - -async function mapMethodsWithConcurrency( - methodSymbols: readonly MethodSymbol[], - token: vscode.CancellationToken, - mapper: (method: MethodSymbol) => Promise, -): Promise { - const results: T[] = new Array(methodSymbols.length); - let nextIndex = 0; - - const runNext = async (): Promise => { - while (!token.isCancellationRequested) { - const currentIndex = nextIndex; - nextIndex += 1; - if (currentIndex >= methodSymbols.length) { - return; - } - results[currentIndex] = await mapper(methodSymbols[currentIndex]); - } - }; - - const workerCount = Math.min(CODE_LENS_CONCURRENCY, methodSymbols.length); - await Promise.all(Array.from({length: workerCount}, () => runNext())); - return results; -} - -async function describeLocation( - location: vscode.Location, - cache: Map>, -): Promise { - const symbols = await getDocumentSymbols(location.uri, cache); - const symbolName = findSymbolNameAtRange(symbols, location.range); - if (symbolName) { - return symbolName; - } - return `${path.basename(location.uri.fsPath)}:${location.range.start.line + 1}`; -} - function singleTargetCommand(target: NavigationTarget): vscode.Command { return { title: '', @@ -210,112 +86,181 @@ function implementationLensTitle(targets: readonly NavigationTarget[]): string { } export class OverrideCodeLensProvider implements vscode.CodeLensProvider { + private readonly requests = new Map(); + private cachedLenses = new WeakMap(); + private readonly retries = new Map(); + private refreshTimer: ReturnType | undefined; + private refreshSubscription: vscode.Disposable | undefined; + private readonly onDidChangeEmitter = new vscode.EventEmitter(); public readonly onDidChangeCodeLenses = this.onDidChangeEmitter.event; - public refresh(): void { - this.onDidChangeEmitter.fire(); + public constructor(private readonly getClient: () => LanguageClient) {} + + public connectClient(client: LanguageClient): void { + this.refreshSubscription?.dispose(); + this.refreshSubscription = client.onNotification(REFRESH_METHOD, () => this.refresh()); + this.cachedLenses = new WeakMap(); + this.refresh(); } - public async provideCodeLenses( - document: vscode.TextDocument, - token: vscode.CancellationToken, - ): Promise { - if (document.languageId !== 'python' || codeLensDisabled(document)) { - return []; + public refresh(): void { + for (const request of this.requests.values()) { + request.cancel(); } + this.retries.clear(); + this.scheduleRefresh(); + } - const rootSymbolsResult = - (await vscode.commands.executeCommand< - vscode.DocumentSymbol[] | vscode.SymbolInformation[] - >( - 'vscode.executeDocumentSymbolProvider', - document.uri, - )) ?? []; - if (!hasHierarchicalSymbols(rootSymbolsResult) || rootSymbolsResult.length === 0) { - return []; - } - const rootSymbols = rootSymbolsResult; - if (rootSymbols.length === 0) { - return []; + private scheduleRefresh(): void { + if (this.refreshTimer !== undefined) { + clearTimeout(this.refreshTimer); } + this.refreshTimer = setTimeout(() => { + this.refreshTimer = undefined; + this.onDidChangeEmitter.fire(); + }, REFRESH_DELAY_MS); + } - const symbolCache = new Map>(); - symbolCache.set(document.uri.toString(), Promise.resolve(rootSymbols)); - - const methodSymbols = collectMethodSymbols(rootSymbols); - if (methodSymbols.length > MAX_METHOD_SYMBOLS) { - return []; + public dispose(): void { + for (const request of this.requests.values()) { + request.cancel(); + request.dispose(); } - - const codeLenses = await mapMethodsWithConcurrency( - methodSymbols, - token, - method => this.getCodeLensesForMethod(document, method, symbolCache, token), - ); - return codeLenses.flat(); + this.requests.clear(); + this.cachedLenses = new WeakMap(); + this.retries.clear(); + this.refreshSubscription?.dispose(); + if (this.refreshTimer !== undefined) { + clearTimeout(this.refreshTimer); + } + this.onDidChangeEmitter.dispose(); } - private async getCodeLensesForMethod( + public async provideCodeLenses( document: vscode.TextDocument, - method: MethodSymbol, - symbolCache: Map>, token: vscode.CancellationToken, ): Promise { - if (token.isCancellationRequested) { + if (document.languageId !== 'python' || codeLensDisabled(document)) { return []; } - const [declarationResult, implementationResult] = await Promise.all([ - vscode.commands.executeCommand( - 'vscode.executeDeclarationProvider', - document.uri, - method.position, - ), - vscode.commands.executeCommand( - 'vscode.executeImplementationProvider', - document.uri, - method.position, - ), - ]); - - const declarationTargets = normalizeLocations(declarationResult).filter( - location => !locationsEqual(location, document.uri, method.range), - ); - const implementationTargets = normalizeLocations(implementationResult).filter( - location => !locationsEqual(location, document.uri, method.range), - ); - - const lenses: vscode.CodeLens[] = []; - if (declarationTargets.length > 0) { - const targets = await Promise.all( - declarationTargets.map(async location => ({ - name: await describeLocation(location, symbolCache), - uri: location.uri, - range: location.range, - })), - ); - const command = targets.length === 1 ? singleTargetCommand(targets[0]) : multiTargetCommand(targets); - command.title = overrideLensTitle(targets); - lenses.push(new vscode.CodeLens(method.range, command)); + const key = document.uri.toString(); + this.requests.get(key)?.cancel(); + const cancellation = new vscode.CancellationTokenSource(); + this.requests.set(key, cancellation); + const subscription = token.onCancellationRequested(() => cancellation.cancel()); + if (token.isCancellationRequested) { + cancellation.cancel(); } + token = cancellation.token; + const version = document.version; + try { + const rootSymbolsResult = + (await vscode.commands.executeCommand< + vscode.DocumentSymbol[] | vscode.SymbolInformation[] + >( + 'vscode.executeDocumentSymbolProvider', + document.uri, + )) ?? []; + if (!hasHierarchicalSymbols(rootSymbolsResult) || rootSymbolsResult.length === 0) { + return []; + } + const rootSymbols = rootSymbolsResult; - if (implementationTargets.length > 0) { - const targets = await Promise.all( - implementationTargets.map(async location => ({ - name: await describeLocation(location, symbolCache), - uri: location.uri, - range: location.range, - })), + const methodRanges = collectMethodRanges(rootSymbols); + if (methodRanges.length > MAX_METHOD_SYMBOLS) { + return []; + } + + if ( + token.isCancellationRequested || + document.version !== version || + methodRanges.length === 0 + ) { + return []; + } + const client = this.getClient(); + const result = await client.sendRequest< + {parents: NamedLocation[]; implementations: NamedLocation[]}[] + >( + 'pyrefly/textDocument/overrideCodeLens', + { + textDocument: + client.code2ProtocolConverter.asVersionedTextDocumentIdentifier(document), + refreshSupport: true, + ranges: methodRanges.map(range => + client.code2ProtocolConverter.asRange(range), + ), + }, + token, ); - const command = - targets.length === 1 ? singleTargetCommand(targets[0]) : multiTargetCommand(targets); - command.title = implementationLensTitle(targets); - lenses.push(new vscode.CodeLens(method.range, command)); + if (token.isCancellationRequested || document.version !== version) { + return []; + } + const toTarget = (location: NamedLocation): NavigationTarget => { + const target = client.protocol2CodeConverter.asLocation(location); + return { + name: location.name ?? + `${path.basename(target.uri.fsPath)}:${target.range.start.line + 1}`, + uri: target.uri, + range: target.range, + }; + }; + const lenses = result.flatMap((targets, i) => { + const methodLenses: vscode.CodeLens[] = []; + for (const [locations, title] of [ + [targets.parents, overrideLensTitle], + [targets.implementations, implementationLensTitle], + ] as const) { + if (locations.length === 0) { + continue; + } + const navigationTargets = locations.map(toTarget); + const command = + navigationTargets.length === 1 + ? singleTargetCommand(navigationTargets[0]) + : multiTargetCommand(navigationTargets); + command.title = title(navigationTargets); + methodLenses.push(new vscode.CodeLens(methodRanges[i], command)); + } + return methodLenses; + }); + this.cachedLenses.set(document, {version, lenses}); + this.retries.delete(key); + return lenses; + } catch (error) { + if (token.isCancellationRequested || document.version !== version) { + return []; + } + if ( + error instanceof ResponseError && + [ + LSPErrorCodes.RequestCancelled, + LSPErrorCodes.ServerCancelled, + LSPErrorCodes.ContentModified, + ].includes(error.code) + ) { + const attempts = this.retries.get(key) ?? 0; + if (attempts < MAX_SERVER_RETRIES) { + this.retries.set(key, attempts + 1); + this.scheduleRefresh(); + } + const cached = this.cachedLenses.get(document); + return cached?.version === version ? cached.lenses : []; + } + throw error; + } finally { + subscription.dispose(); + cancellation.dispose(); + if (this.requests.get(key) === cancellation) { + this.requests.delete(key); + } } - - return lenses; } } diff --git a/pyrefly/lib/lsp/non_wasm/code_lens.rs b/pyrefly/lib/lsp/non_wasm/code_lens.rs index 594358e72b..aa05fe6595 100644 --- a/pyrefly/lib/lsp/non_wasm/code_lens.rs +++ b/pyrefly/lib/lsp/non_wasm/code_lens.rs @@ -7,8 +7,10 @@ use lsp_types::CodeLens; use lsp_types::Command; +use lsp_types::Location; use lsp_types::Range; use lsp_types::Uri; +use lsp_types::VersionedTextDocumentIdentifier; use pyrefly_build::handle::Handle; use pyrefly_python::ast::Ast; use ruff_python_ast::Expr; @@ -16,10 +18,58 @@ use ruff_python_ast::ExprAttribute; use ruff_python_ast::Stmt; use ruff_python_ast::StmtClassDef; use ruff_text_size::TextRange; +use serde::Deserialize; +use serde::Serialize; use serde_json::Value; use crate::state::state::Transaction; +/// A document batch for override lenses. Ranges use the specified editor version. +#[derive(Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct OverrideCodeLensParams { + pub text_document: VersionedTextDocumentIdentifier, + pub ranges: Vec, + /// The client handles `pyrefly/overrideCodeLens/refresh` notifications. + #[serde(default)] + pub refresh_support: bool, +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct OverrideCodeLensTarget { + #[serde(flatten)] + pub location: Location, + pub name: String, +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct OverrideCodeLensTargets { + pub parents: Vec, + pub implementations: Vec, +} + +/// Invalidates override lenses after the server updates its analysis. +pub enum OverrideCodeLensRefreshNotification {} + +impl lsp_types::Notification for OverrideCodeLensRefreshNotification { + type Params = Value; + const METHOD: lsp_types::LspNotificationMethod<'static> = + lsp_types::LspNotificationMethod::Custom("pyrefly/overrideCodeLens/refresh"); + const MESSAGE_DIRECTION: lsp_types::MessageDirection = + lsp_types::MessageDirection::ServerToClient; +} + +pub enum OverrideCodeLensRequest {} + +impl lsp_types::Request for OverrideCodeLensRequest { + type Params = OverrideCodeLensParams; + type Result = Vec; + const METHOD: lsp_types::LspRequestMethod<'static> = + lsp_types::LspRequestMethod::Custom("pyrefly/textDocument/overrideCodeLens"); + const MESSAGE_DIRECTION: lsp_types::MessageDirection = + lsp_types::MessageDirection::ClientToServer; +} + #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum CodeLensKind { Run, diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 1e9eee5457..03ed5151c1 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -299,6 +299,11 @@ use crate::lsp::non_wasm::call_hierarchy::find_function_at_position_in_ast; use crate::lsp::non_wasm::call_hierarchy::prepare_call_hierarchy_item; use crate::lsp::non_wasm::call_hierarchy::transform_incoming_calls; use crate::lsp::non_wasm::call_hierarchy::transform_outgoing_calls; +use crate::lsp::non_wasm::code_lens::OverrideCodeLensParams; +use crate::lsp::non_wasm::code_lens::OverrideCodeLensRefreshNotification; +use crate::lsp::non_wasm::code_lens::OverrideCodeLensRequest; +use crate::lsp::non_wasm::code_lens::OverrideCodeLensTarget; +use crate::lsp::non_wasm::code_lens::OverrideCodeLensTargets; use crate::lsp::non_wasm::code_lens::runnable_lsp_code_lens; use crate::lsp::non_wasm::convert_module_package::convert_module_package_code_actions; use crate::lsp::non_wasm::document_symbols::flatten_to_symbol_information; @@ -1217,6 +1222,7 @@ pub struct Server { /// These registrations are additive and never removed, so this set only grows. watched_exact_paths: Mutex>, version_info: Mutex>, + override_code_lens_refresh_support: AtomicBool, id: Uuid, /// The surface/entrypoint for the language server (`--from` CLI arg) surface: Option, @@ -2232,7 +2238,13 @@ impl Server { .collect(); let ruff_notebook = params.notebook_document.to_ruff_notebook(&cell_contents)?; - let lsp_notebook = LspNotebook::new(ruff_notebook, notebook_document); + let cell_versions = params + .cell_text_documents + .iter() + .map(|doc| (doc.uri.clone(), doc.version)) + .collect(); + let lsp_notebook = + LspNotebook::new(ruff_notebook, notebook_document, cell_versions); let notebook_path = url.to_file_path().unwrap_or_else(|_| { self.unsaved_file_tracker .ensure_path_for_open(&url, "jupyter") @@ -2444,6 +2456,7 @@ impl Server { Some(&self.lsp_thread_pool), ); info!("Handling non-canceled request {} ({})", x.method, &x.id); + let mut queued_override_code_lens = false; if let Some(params) = as_request::(&x) { if let Some(params) = self .extract_request_params_or_send_err_response::( @@ -2686,6 +2699,27 @@ impl Server { }; self.send_response(new_response(x.id, Ok(response))); } + } else if let Some(params) = as_request::(&x) { + if let Some(params) = self + .extract_request_params_or_send_err_response::( + params, &x.id, + ) + { + match self.async_override_code_lens( + x.id.clone(), + params, + telemetry_event.activity_key.clone(), + ) { + Ok(()) => queued_override_code_lens = true, + Err(reason) => { + telemetry_event.set_empty_response_reason(reason); + self.send_response(new_response( + x.id, + Ok(Vec::::new()), + )); + } + } + } } else if let Some(params) = as_request::(&x) { if let Some(params) = self .extract_request_params_or_send_err_response::( @@ -2982,9 +3016,11 @@ impl Server { )); info!("Unhandled request: {x:?}"); } - self.cancellation_handles - .lock() - .remove(&request_id_for_cancel); + if !queued_override_code_lens { + self.cancellation_handles + .lock() + .remove(&request_id_for_cancel); + } ide_transaction_manager.save(transaction, telemetry_event); } } @@ -3086,6 +3122,7 @@ impl Server { watched_patterns: Mutex::new(SmallSet::new()), watched_exact_paths: Mutex::new(SmallSet::new()), version_info: Mutex::new(HashMap::new()), + override_code_lens_refresh_support: AtomicBool::new(false), id: Uuid::new_v4(), surface, agent_session_id, @@ -3547,6 +3584,18 @@ impl Server { { self.send_request::(()); } + if !handles.is_empty() + && self + .override_code_lens_refresh_support + .load(Ordering::Relaxed) + { + self.connection + .send(Message::Notification(new_notification::< + OverrideCodeLensRefreshNotification, + >(serde_json::json!( + {} + )))); + } } /// Validate open files and send errors to the LSP. In the case of an ongoing recheck @@ -4452,6 +4501,7 @@ impl Server { }; let mut notebook_document = original_notebook.notebook_document().clone(); + let mut cell_versions = original_notebook.cell_versions().clone(); let mut cell_content_map: HashMap = HashMap::new(); // Changed metadata if let Some(metadata) = ¶ms.change.metadata { @@ -4500,6 +4550,7 @@ impl Server { // Set contents for new cells if let Some(opened_cells) = &structure.did_open { for opened_cell in opened_cells { + cell_versions.insert(opened_cell.uri.clone(), opened_cell.version); cell_content_map.insert(opened_cell.uri.clone(), opened_cell.text.clone()); self.open_notebook_cells .write() @@ -4535,6 +4586,7 @@ impl Server { .filter_map(|v| serde_json::from_value(v.clone()).ok()) .collect(); let new_text = apply_change_events(original_text, content_changes); + cell_versions.insert(cell_uri.clone(), text_change.document.version); cell_content_map.insert(cell_uri, new_text); } } @@ -4544,7 +4596,11 @@ impl Server { .clone() .to_ruff_notebook(&cell_content_map)?; - let new_notebook = Arc::new(LspNotebook::new(ruff_notebook, notebook_document)); + let new_notebook = Arc::new(LspNotebook::new( + ruff_notebook, + notebook_document, + cell_versions, + )); *original = Arc::new(LspFile::Notebook(new_notebook)); drop(lock); // Update version_info only after the mutation has fully succeeded, so @@ -5994,6 +6050,181 @@ impl Server { Ok(Some(res)) } + /// Returns the text document version, including versions of individual notebook cells. + fn text_document_version(&self, uri: &Uri) -> Option { + let path = self.path_for_uri(uri)?; + if self.open_notebook_cells.read().contains_key(uri) { + let files = self.open_files.read(); + let LspFile::Notebook(notebook) = files.get(&path)?.as_ref() else { + return None; + }; + notebook.cell_versions().get(uri).copied() + } else { + self.version_info.lock().get(&path).copied() + } + } + + /// Queues one cancellable search for all method lenses in a document. + fn async_override_code_lens( + &self, + request_id: RequestId, + params: OverrideCodeLensParams, + activity_key: Option, + ) -> Result<(), EmptyResponseReason> { + let uri = ¶ms.text_document.text_document_identifier.uri; + let handle = self.make_handle_if_enabled(uri, Some(CodeLensRequest::METHOD.as_str()))?; + if params.refresh_support { + self.override_code_lens_refresh_support + .store(true, Ordering::Relaxed); + } + // Register before enqueueing so cancellation also stops work waiting in the queue. + let queued_cancellation = CancellationHandle::default(); + self.cancellation_handles + .lock() + .insert(request_id.clone(), queued_cancellation.clone()); + self.find_reference_queue.queue_task( + TelemetryEventKind::FindFromDefinition, + Box::new(move |server, _, telemetry_event| { + telemetry_event.set_activity_key(activity_key); + let mut transaction = server.state.cancellable_transaction(); + let result = (|| { + { + // Replace the queued token atomically with the running transaction token. + let mut handles = server.cancellation_handles.lock(); + if queued_cancellation.is_cancelled() { + return Err(Cancelled); + } + handles.insert(request_id.clone(), transaction.get_cancellation_handle()); + } + if server + .text_document_version(¶ms.text_document.text_document_identifier.uri) + != Some(params.text_document.version) + { + return Err(Cancelled); + } + server.validate_in_memory_for_transaction( + transaction.as_mut(), + telemetry_event, + None, + ); + transaction.run(&[handle.dupe()], Require::Everything, None)?; + let info = transaction + .as_ref() + .get_module_info(&handle) + .ok_or(Cancelled)?; + let ranges = params + .ranges + .iter() + .map(|range| { + info.from_lsp_range( + *range, + server.maybe_get_code_cell_index( + ¶ms.text_document.text_document_identifier.uri, + ), + ) + }) + .collect::>(); + let targets = transaction.override_code_lens_targets(&handle, &ranges)?; + if transaction.get_cancellation_handle().is_cancelled() + || server.text_document_version( + ¶ms.text_document.text_document_identifier.uri, + ) != Some(params.text_document.version) + { + return Err(Cancelled); + } + let notebooks = server.snapshot_open_notebooks(); + let mut names = SmallMap::>::new(); + let mut to_locations = |targets: Vec| { + targets + .into_iter() + .filter_map(|target| { + let mut uri = module_info_to_uri( + &target.module, + server.path_remapper.as_ref(), + )?; + if let Some(cell) = + target.module.to_cell_for_lsp(target.range.start()) + && let Some(path) = to_real_path(target.module.path()) + && let Some(notebook) = notebooks.get(&path) + && let Some(cell_uri) = notebook.get_code_cell_url(cell) + { + uri = cell_uri.clone(); + } + let module_names = + names.entry(target.module.path().dupe()).or_insert_with(|| { + let target_handle = Handle::new( + target.module.name(), + target.module.path().dupe(), + *handle.sys_info(), + ); + transaction + .as_ref() + .get_exports_data(&target_handle) + .symbols() + .map(|symbols| { + symbols + .iter() + .map(|(symbol, parent)| { + let name = target + .module + .code_at(symbol.name.range()); + let name = match parent { + Some(parent) => format!( + "{}.{}", + target + .module + .code_at(parent.name.range()), + name + ), + None => name.to_owned(), + }; + (symbol.name.range(), name) + }) + .collect() + }) + .unwrap_or_default() + }); + let name = + module_names.get(&target.range).cloned().unwrap_or_else(|| { + target.module.code_at(target.range).to_owned() + }); + Some(OverrideCodeLensTarget { + location: Location { + uri, + range: target.module.to_lsp_range(target.range), + }, + name, + }) + }) + .collect() + }; + Ok(targets + .into_iter() + .map(|targets| OverrideCodeLensTargets { + parents: to_locations(targets.parents), + implementations: to_locations(targets.implementations), + }) + .collect::>()) + })(); + server.cancellation_handles.lock().remove(&request_id); + let response = match result { + Ok(targets) => new_response(request_id, Ok(targets)), + Err(Cancelled) => { + telemetry_event.canceled = true; + Response::new_err( + request_id, + ErrorCode::RequestCanceled as i32, + "Override CodeLens request is canceled or its document version changed" + .to_owned(), + ) + } + }; + server.connection.send(Message::Response(response)); + }), + ); + Ok(()) + } + fn code_lens( &self, transaction: &Transaction<'_>, diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 774dc60bc8..9c9dd6ea5e 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -7,6 +7,7 @@ use std::cmp::Ordering; use std::cmp::Reverse; +use std::collections::HashMap; use std::collections::HashSet; use std::iter; use std::path::PathBuf; @@ -3374,41 +3375,44 @@ impl<'a> Transaction<'a> { }) } - /// Finds overridden parent declarations for a method or class attribute definition. - /// - /// The answers index records override edges from a child declaration range to the - /// declaration ranges it overrides in the MRO. We use that mapping to make - /// go-to-declaration on an overriding member jump to the parent declaration instead - /// of back to the local implementation. - fn parent_declarations_for_definition( + /// Finds ancestor declarations directly from class fields and the MRO. + /// Each range belongs to the module snapshot stored on its declaring class. + pub fn parent_declarations( &self, handle: &Handle, - definition: &TextRangeWithModule, - module_cache: &mut SmallMap>, + range: TextRange, ) -> Vec { - let Some(index) = solutions_index_for_module(self, handle, &definition.module) else { + let Some(answers) = self.get_answers(handle) else { return Vec::new(); }; - let Some(parent_methods) = index - .lock() - .parent_methods_map - .get(&definition.range) - .cloned() - else { + let Some(class_key) = answers.bindings().enclosing_class(range) else { return Vec::new(); }; - - let mut seen = HashSet::new(); - let mut declarations = Vec::new(); - for (parent_module_path, parent_range) in parent_methods { - let Some(module) = module_info_for_path(self, module_cache, &parent_module_path) else { - continue; - }; - if seen.insert((module.path().dupe(), parent_range)) { - declarations.push(TextRangeWithModule::new(module, parent_range)); + self.ad_hoc_solve(handle, "parent_declarations", |solver| { + let class = solver.get_idx(class_key).0.as_ref()?; + let fields = solver.get_class_fields(class)?; + let name = fields + .names() + .find(|name| fields.field_decl_range(name) == Some(range))?; + if Ast::is_mangled_attr(name) { + return Some(Vec::new()); } - } - declarations + let mut seen = HashSet::new(); + Some( + solver + .get_mro_for_class(class) + .ancestors(solver.stdlib) + .filter_map(|ancestor| { + let class = ancestor.class_object(); + let range = solver.get_class_fields(class)?.field_decl_range(name)?; + seen.insert((class.module().path().dupe(), range)) + .then(|| TextRangeWithModule::new(class.module().dupe(), range)) + }) + .collect(), + ) + }) + .flatten() + .unwrap_or_default() } pub fn goto_declaration( @@ -3427,26 +3431,9 @@ impl<'a> Transaction<'a> { ..Default::default() }, )?; - let mut module_cache = SmallMap::new(); - let mut seen = HashSet::new(); - let mut results = Vec::new(); - for definition in definitions { - let base = - TextRangeWithModule::new(definition.module.dupe(), definition.definition_range); - let declarations = - self.parent_declarations_for_definition(handle, &base, &mut module_cache); - let declarations = if declarations.is_empty() { - vec![base] - } else { - declarations - }; - for declaration in declarations { - if seen.insert((declaration.module.path().dupe(), declaration.range)) { - results.push(declaration); - } - } - } - Ok(results) + Ok(definitions + .into_vec() + .into_map(|item| TextRangeWithModule::new(item.module, item.definition_range))) } /// Where a function-valued type was defined. @@ -5413,58 +5400,6 @@ fn find_child_implementations_impl( child_implementations } -fn module_info_for_path( - transaction: &Transaction<'_>, - cache: &mut SmallMap>, - module_path: &ModulePath, -) -> Option { - if let Some(module) = cache.get(module_path) { - return module.dupe(); - } - - let counterpart_path = match module_path.details() { - ModulePathDetails::Memory(path) => Some(ModulePath::filesystem((**path).clone())), - ModulePathDetails::FileSystem(path) => Some(ModulePath::memory(path.clone().to_path_buf())), - _ => None, - }; - - let module = transaction.handles().into_iter().find_map(|handle| { - let path = handle.path(); - ((path == module_path) - || counterpart_path - .as_ref() - .is_some_and(|counterpart| path == counterpart)) - .then(|| transaction.get_module_info(&handle)) - .flatten() - }); - cache.insert(module_path.dupe(), module.dupe()); - if let Some(counterpart_path) = counterpart_path { - cache.insert(counterpart_path, module.dupe()); - } - module -} - -fn solutions_index_for_module( - transaction: &Transaction<'_>, - handle: &Handle, - module: &Module, -) -> Option>> { - let counterpart_path = match module.path().details() { - ModulePathDetails::Memory(path) => Some(ModulePath::filesystem((**path).clone())), - ModulePathDetails::FileSystem(path) => Some(ModulePath::memory(path.clone().to_path_buf())), - _ => None, - }; - std::iter::once(module.path().dupe()) - .chain(counterpart_path) - .find_map(|path| { - self::Transaction::get_solutions( - transaction, - &Handle::new(module.name(), path, handle.sys_info().dupe()), - ) - .and_then(|solutions| solutions.get_index()) - }) -} - fn compute_transitive_rdeps_for_definition_impl( transaction: &mut T, sys_info: SysInfo, @@ -5478,11 +5413,12 @@ fn compute_transitive_rdeps_for_definition_impl( sys_info, ); let mut rdeps = transaction.transitive_rdeps(handle_of_filesystem_counterpart.dupe()); - rdeps.insert(Handle::new( + // Imports can also resolve directly to an in-memory module. + rdeps.extend(transaction.transitive_rdeps(Handle::new( definition.module.name(), definition.module.path().dupe(), sys_info, - )); + ))); rdeps } _ => { @@ -5687,7 +5623,101 @@ impl<'a> Transaction<'a> { } } +/// Ancestor declarations and child implementations for one method. +pub struct OverrideTargets { + pub parents: Vec, + pub implementations: Vec, +} + impl<'a> CancellableTransaction<'a> { + /// Batches override navigation for one document into one reverse-dependency search. + pub fn override_code_lens_targets( + &mut self, + handle: &Handle, + ranges: &[TextRange], + ) -> Result, Cancelled> { + let cancellation = self.get_cancellation_handle(); + let mut results = Vec::with_capacity(ranges.len()); + for range in ranges { + if cancellation.is_cancelled() { + return Err(Cancelled); + } + results.push(OverrideTargets { + parents: self.as_ref().parent_declarations(handle, *range), + implementations: Vec::new(), + }); + } + if ranges.is_empty() { + return Ok(results); + } + let Some(module) = self.as_ref().get_module_info(handle) else { + return Ok(results); + }; + let definition = + TextRangeWithModule::new(module.dupe(), TextRange::empty(TextSize::new(0))); + let candidates = + compute_transitive_rdeps_for_definition_impl(self, *handle.sys_info(), &definition)?; + let mut methods = HashMap::new(); + let filesystem_handle = match module.path().details() { + ModulePathDetails::Memory(path) => Some(Handle::new( + module.name(), + ModulePath::filesystem((**path).clone()), + *handle.sys_info(), + )), + _ => None, + }; + for (i, range) in ranges.iter().enumerate() { + methods.insert((module.path().dupe(), *range), i); + if let Some(filesystem_handle) = &filesystem_handle { + let patched = patch_definition_for_handle_impl( + self, + filesystem_handle, + &TextRangeWithModule::new(module.dupe(), *range), + ); + methods.insert((patched.module.path().dupe(), patched.range), i); + } + } + for candidate in candidates { + if cancellation.is_cancelled() { + return Err(Cancelled); + } + let Some(index) = self.solutions_index(&candidate) else { + continue; + }; + let index = index.lock(); + if index.parent_methods_map.is_empty() { + continue; + } + let Some(child_module) = self.as_ref().get_module_info(&candidate) else { + continue; + }; + for (child_range, parents) in &index.parent_methods_map { + if cancellation.is_cancelled() { + return Err(Cancelled); + } + for (parent_path, parent_range) in parents { + if let Some(i) = methods.get(&(parent_path.dupe(), *parent_range)) { + results[*i] + .implementations + .push(TextRangeWithModule::new(child_module.dupe(), *child_range)); + } + } + } + } + for targets in &mut results { + let children = &mut targets.implementations; + children.sort_by_key(|child| { + ( + child.module.path().dupe(), + child.range.start(), + child.range.end(), + ) + }); + children.dedup_by_key(|child| (child.module.path().dupe(), child.range)); + } + Ok(results) + } + /// Processes each transitive reverse dependency for a given definition location. /// /// This is a common pattern in workspace-wide references-related features. Candidates are diff --git a/pyrefly/lib/state/notebook.rs b/pyrefly/lib/state/notebook.rs index 3924bd5615..97e08b31aa 100644 --- a/pyrefly/lib/state/notebook.rs +++ b/pyrefly/lib/state/notebook.rs @@ -21,24 +21,36 @@ pub struct LspNotebook { // Notebook cells have Urls of unspecified format cell_url_to_index: SmallMap, cell_index_to_url: Vec, + cell_versions: SmallMap, } impl LspNotebook { - pub fn new(ruff_notebook: Notebook, notebook_document: NotebookDocument) -> Self { + pub fn new( + ruff_notebook: Notebook, + notebook_document: NotebookDocument, + mut cell_versions: SmallMap, + ) -> Self { let mut cell_url_to_index = SmallMap::new(); let mut cell_index_to_url = Vec::new(); for (idx, cell) in notebook_document.cells.iter().enumerate() { cell_url_to_index.insert(cell.document.clone(), idx); cell_index_to_url.push(cell.document.clone()); } + cell_versions.retain(|uri, _| cell_url_to_index.contains_key(uri)); Self { ruff_notebook: Arc::new(ruff_notebook), notebook_document, cell_url_to_index, cell_index_to_url, + cell_versions, } } + /// Cell text versions are independent of the notebook version. + pub fn cell_versions(&self) -> &SmallMap { + &self.cell_versions + } + pub fn notebook_document(&self) -> &NotebookDocument { &self.notebook_document } @@ -140,7 +152,7 @@ mod tests { .clone() .to_ruff_notebook(&cell_content) .unwrap(); - let lsp_notebook = LspNotebook::new(ruff_notebook, notebook_doc); + let lsp_notebook = LspNotebook::new(ruff_notebook, notebook_doc, SmallMap::new()); // First code cell (all-cells index 0) → code-cell index 0 assert_eq!(lsp_notebook.get_code_cell_index(&cell0_url), Some(0));