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 be948979ac..dd05604c71 100644 --- a/lsp/src/extension.ts +++ b/lsp/src/extension.ts @@ -36,12 +36,18 @@ 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'; import {describeError, resolveExecutable} from './lsp-path'; let client: LanguageClient; +let overrideCodeLensProvider: OverrideCodeLensProvider; let outputChannel: vscode.OutputChannel; let traceOutputChannel: vscode.OutputChannel; let inferOutputChannel: vscode.OutputChannel; @@ -96,6 +102,7 @@ async function replaceClient(): Promise { launchSpec, clientOptions, ); + overrideCodeLensProvider.connectClient(client); try { await client.start(); logServerVersion(); @@ -252,6 +259,9 @@ export async function activate( } pythonEnv = new PythonEnvironment(context); + overrideCodeLensProvider = new OverrideCodeLensProvider(() => client); + context.subscriptions.push(overrideCodeLensProvider); + const scheduleCodeLensRefresh = () => overrideCodeLensProvider.refresh(); launchSpec = await resolveExecutable( context.extensionUri, @@ -343,12 +353,26 @@ export async function activate( serverOptions, clientOptions, ); + overrideCodeLensProvider.connectClient(client); if (supportsHoverVerbosity) { 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 () => { + scheduleCodeLensRefresh(); await updateStatusBar(client); }), ); @@ -386,6 +410,7 @@ export async function activate( client.sendNotification(DidChangeConfigurationNotification.type, { settings: {}, }); + scheduleCodeLensRefresh(); } if ( event.affectsConfiguration('pyrefly.lspPath') || @@ -401,6 +426,14 @@ export async function activate( }), ); + context.subscriptions.push( + workspace.onDidChangeTextDocument(event => { + if (event.document.languageId === 'python') { + scheduleCodeLensRefresh(); + } + }), + ); + context.subscriptions.push( vscode.commands.registerCommand('pyrefly.restartClient', async () => { // Clear the output channel but don't dispose it @@ -435,6 +468,7 @@ export async function activate( ); } }); + scheduleCodeLensRefresh(); }), ); @@ -510,6 +544,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 +582,7 @@ export async function activate( } }); logServerVersion(); + scheduleCodeLensRefresh(); await updateStatusBar(client); const statusBarItem = getStatusBarItem(); diff --git a/lsp/src/overrideCodeLens.ts b/lsp/src/overrideCodeLens.ts new file mode 100644 index 0000000000..4cdb17ef30 --- /dev/null +++ b/lsp/src/overrideCodeLens.ts @@ -0,0 +1,298 @@ +/** + * 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'; +import {LanguageClient, Location, LSPErrorCodes, ResponseError} from 'vscode-languageclient/node'; + +const MAX_METHOD_SYMBOLS = 200; +const REFRESH_DELAY_MS = 150; +const MAX_SERVER_RETRIES = 3; +const REFRESH_METHOD = 'pyrefly/overrideCodeLens/refresh'; + +type NamedLocation = Location & {name?: string}; + +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 hasHierarchicalSymbols( + symbols: vscode.DocumentSymbol[] | vscode.SymbolInformation[] | undefined, +): symbols is vscode.DocumentSymbol[] { + return symbols !== undefined && (symbols.length === 0 || 'children' in symbols[0]); +} + +function collectMethodRanges( + symbols: readonly vscode.DocumentSymbol[], + inClass = false, +): vscode.Range[] { + const methods: vscode.Range[] = []; + for (const symbol of symbols) { + if (symbol.kind === vscode.SymbolKind.Method && inClass) { + methods.push(symbol.selectionRange); + } else { + methods.push( + ...collectMethodRanges( + symbol.children, + inClass || symbol.kind === vscode.SymbolKind.Class, + ), + ); + } + } + return methods; +} + +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 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 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 refresh(): void { + for (const request of this.requests.values()) { + request.cancel(); + } + this.retries.clear(); + this.scheduleRefresh(); + } + + private scheduleRefresh(): void { + if (this.refreshTimer !== undefined) { + clearTimeout(this.refreshTimer); + } + this.refreshTimer = setTimeout(() => { + this.refreshTimer = undefined; + this.onDidChangeEmitter.fire(); + }, REFRESH_DELAY_MS); + } + + public dispose(): void { + for (const request of this.requests.values()) { + request.cancel(); + request.dispose(); + } + this.requests.clear(); + this.cachedLenses = new WeakMap(); + this.retries.clear(); + this.refreshSubscription?.dispose(); + if (this.refreshTimer !== undefined) { + clearTimeout(this.refreshTimer); + } + this.onDidChangeEmitter.dispose(); + } + + public async provideCodeLenses( + document: vscode.TextDocument, + token: vscode.CancellationToken, + ): Promise { + if (document.languageId !== 'python' || codeLensDisabled(document)) { + return []; + } + + 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; + + 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, + ); + 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); + } + } + } +} + +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/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/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 8088d344cc..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,6 +3375,46 @@ impl<'a> Transaction<'a> { }) } + /// 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, + range: TextRange, + ) -> Vec { + let Some(answers) = self.get_answers(handle) else { + return Vec::new(); + }; + let Some(class_key) = answers.bindings().enclosing_class(range) else { + return Vec::new(); + }; + 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()); + } + 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( &self, handle: &Handle, @@ -3390,7 +3431,6 @@ impl<'a> Transaction<'a> { ..Default::default() }, )?; - Ok(definitions .into_vec() .into_map(|item| TextRangeWithModule::new(item.module, item.definition_range))) @@ -5373,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 } _ => { @@ -5582,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)); diff --git a/pyrefly/lib/test/lsp/definition.rs b/pyrefly/lib/test/lsp/definition.rs index 95e70b9732..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; @@ -77,6 +79,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 +512,142 @@ Definition Result: ); } +#[test] +fn overriding_method_declaration_stays_at_definition_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: +7 | def f(self) -> int: + ^ +"# + .trim(), + report.trim(), + ); +} + +#[test] +fn overriding_class_attribute_declaration_stays_at_definition_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: +6 | value = 2 + ^^^^^ +"# + .trim(), + report.trim(), + ); +} + +#[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(); +}