diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index d1be83be33..8b8759387b 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -47,6 +47,7 @@ use lsp_types::CodeActionTriggerKind; use lsp_types::CodeLens; use lsp_types::CodeLensOptions; use lsp_types::CodeLensParams; +use lsp_types::Command; use lsp_types::CompletionItem; use lsp_types::CompletionList; use lsp_types::CompletionOptions; @@ -1101,6 +1102,12 @@ struct InitializeResult { server_info: Option, } +#[derive(Debug, Clone)] +struct CodeLensTarget { + range: Range, + definition: FindDefinitionItemWithDocstring, +} + pub fn initialize_finish( sender: &Sender, reader: &mut MessageReader, @@ -2251,18 +2258,6 @@ 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, - ) - { - self.set_file_stats(params.text_document.uri.clone(), telemetry_event); - self.send_response(new_response( - x.id, - Ok(self.code_lens(&transaction, params).unwrap_or_default()), - )); - } } else if let Some(params) = as_request::(&x) { if let Some(params) = self .extract_request_params_or_send_err_response::( @@ -2309,6 +2304,23 @@ 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, + ) + { + self.set_file_stats(params.text_document.uri.clone(), telemetry_event); + if let Err(reason) = self.code_lens( + x.id.clone(), + &transaction, + params, + telemetry_event.activity_key.clone(), + ) { + self.send_response(new_response(x.id, Ok(None::<()>))); + telemetry_event.set_empty_response_reason(reason); + } + } } else if let Some(params) = as_request::(&x) { if let Some(params) = self .extract_request_params_or_send_err_response::( @@ -4633,6 +4645,64 @@ impl Server { Ok((!actions.is_empty()).then_some(actions)) } + fn code_lens_targets( + &self, + transaction: &Transaction<'_>, + handle: &Handle, + uri: &Url, + ) -> Result, EmptyResponseReason> { + fn recurse_symbols<'a>(symbols: &'a [DocumentSymbol], out: &mut Vec<&'a DocumentSymbol>) { + for symbol in symbols { + if matches!( + symbol.kind, + SymbolKind::CLASS | SymbolKind::FUNCTION | SymbolKind::METHOD + ) { + out.push(symbol); + } + if let Some(children) = symbol.children.as_deref() { + recurse_symbols(children, out); + } + } + } + + let module_info = transaction + .get_module_info(handle) + .ok_or(EmptyResponseReason::ModuleInfoNotFound)?; + let symbols = transaction + .symbols(handle, None) + .ok_or(EmptyResponseReason::ModuleInfoNotFound)?; + let mut symbol_defs = Vec::new(); + recurse_symbols(&symbols, &mut symbol_defs); + + let mut seen = SmallSet::new(); + let mut targets = Vec::new(); + for symbol in symbol_defs { + let position = self.from_lsp_position(uri, &module_info, symbol.selection_range.start); + let definition = match transaction.find_definition( + handle, + position, + FindPreference { + import_behavior: ImportBehavior::StopAtRenamedImports, + ..Default::default() + }, + ) { + Ok(definitions) => definitions.into_vec().swap_remove(0), + Err(_) => { + continue; + } + }; + let key = (definition.module.path().dupe(), definition.definition_range); + if !seen.insert(key) { + continue; + } + targets.push(CodeLensTarget { + range: symbol.selection_range, + definition, + }); + } + Ok(targets) + } + fn document_highlight( &self, transaction: &Transaction<'_>, @@ -4884,6 +4954,150 @@ impl Server { ) } + fn code_lens<'a>( + &'a self, + request_id: RequestId, + transaction: &Transaction<'a>, + params: CodeLensParams, + activity_key: Option, + ) -> Result<(), EmptyResponseReason> { + let uri = ¶ms.text_document.uri; + if self.open_notebook_cells.read().contains_key(uri) { + self.send_response(new_response(request_id, Ok(Some(Vec::::new())))); + return Ok(()); + } + let handle = self.make_handle_if_enabled(uri, Some(CodeLensRequest::METHOD))?; + let runnable_lenses = self.runnable_code_lenses(transaction, &handle, ¶ms)?; + let targets = if self.indexing_mode == IndexingMode::None { + Vec::new() + } else { + self.code_lens_targets(transaction, &handle, uri)? + }; + if targets.is_empty() { + self.send_response(new_response(request_id, Ok(Some(runnable_lenses)))); + return Ok(()); + } + + let path_remapper = self.path_remapper.clone(); + let source_uri = uri.clone(); + self.find_reference_queue.queue_task( + TelemetryEventKind::FindFromDefinition, + Box::new(move |server, _telemetry, telemetry_event| { + telemetry_event.set_activity_key(activity_key); + let mut transaction = server.state.cancellable_transaction(); + server + .cancellation_handles + .lock() + .insert(request_id.clone(), transaction.get_cancellation_handle()); + server.validate_in_memory_for_transaction( + transaction.as_mut(), + telemetry_event, + None, + ); + + let mut lenses = runnable_lenses; + for target in targets { + let local_results = match transaction.find_global_references_from_definition( + *handle.sys_info(), + target.definition.metadata, + TextRangeWithModule::new( + target.definition.module.clone(), + target.definition.definition_range, + ), + false, + ) { + Ok(results) => results, + Err(Cancelled) => { + let message = format!("Request {request_id} is canceled"); + info!("{message}"); + server.connection.send(Message::Response(Response::new_err( + request_id, + ErrorCode::RequestCanceled as i32, + message, + ))); + return; + } + }; + + let mut locations = Vec::new(); + for (info, ranges) in local_results { + if let Some(uri) = module_info_to_uri(&info, path_remapper.as_ref()) { + for range in ranges { + locations.push(Location { + uri: uri.clone(), + range: info.to_lsp_range(range), + }); + } + } + } + + let reference_count = locations.len(); + let title = if reference_count == 1 { + "1 reference".to_owned() + } else { + format!("{reference_count} references") + }; + lenses.push(CodeLens { + range: target.range, + command: Some(Command { + title, + command: "editor.action.showReferences".to_owned(), + arguments: Some(vec![ + serde_json::to_value(&source_uri) + .expect("URI should serialize for code lens"), + serde_json::to_value(target.range.start) + .expect("Position should serialize for code lens"), + serde_json::to_value(&locations) + .expect("Locations should serialize for code lens"), + ]), + }), + data: None, + }); + } + + server.cancellation_handles.lock().remove(&request_id); + server.connection.send(Message::Response(new_response( + request_id, + Ok(Some(lenses)), + ))); + }), + ); + Ok(()) + } + + fn runnable_code_lenses( + &self, + transaction: &Transaction<'_>, + handle: &Handle, + params: &CodeLensParams, + ) -> Result, EmptyResponseReason> { + let uri = ¶ms.text_document.uri; + let path = self + .path_for_uri(uri) + .ok_or(EmptyResponseReason::NoFilePath)?; + let runnable_code_lens = self + .workspaces + .get_with(path.clone(), |(_, workspace)| workspace.runnable_code_lens); + let maybe_cell_idx = self.maybe_get_cell_index(uri); + let info = transaction + .get_module_info(handle) + .ok_or(EmptyResponseReason::ModuleInfoNotFound)?; + let entries = transaction + .runnable_code_lens_entries(handle, uri, runnable_code_lens) + .ok_or(EmptyResponseReason::ModuleInfoNotFound)?; + let cwd = self.runnable_code_lens_cwd(&path); + + let mut lenses = Vec::new(); + for entry in entries { + if info.to_cell_for_lsp(entry.range.start()) != maybe_cell_idx { + continue; + } + let range = info.to_lsp_range(entry.range); + lenses.push(runnable_lsp_code_lens(uri, range, entry, cwd.as_deref())); + } + Ok(lenses) + } + fn rename<'a>( &'a self, request_id: RequestId, @@ -5048,36 +5262,6 @@ impl Server { Ok(Some(res)) } - fn code_lens( - &self, - transaction: &Transaction<'_>, - params: CodeLensParams, - ) -> Option> { - let uri = ¶ms.text_document.uri; - let path = self.path_for_uri(uri)?; - let runnable_code_lens = self - .workspaces - .get_with(path.clone(), |(_, workspace)| workspace.runnable_code_lens); - let maybe_cell_idx = self.maybe_get_cell_index(uri); - let handle = self - .make_handle_if_enabled(uri, Some(CodeLensRequest::METHOD)) - .ok()?; - let info = transaction.get_module_info(&handle)?; - let entries = transaction.runnable_code_lens_entries(&handle, uri, runnable_code_lens)?; - let cwd = self.runnable_code_lens_cwd(&path); - - let mut lenses = Vec::new(); - for entry in entries { - if info.to_cell_for_lsp(entry.range.start()) != maybe_cell_idx { - continue; - } - let range = info.to_lsp_range(entry.range); - lenses.push(runnable_lsp_code_lens(uri, range, entry, cwd.as_deref())); - } - - Some(lenses) - } - fn semantic_tokens_full( &self, transaction: &Transaction<'_>, diff --git a/pyrefly/lib/lsp/non_wasm/workspace.rs b/pyrefly/lib/lsp/non_wasm/workspace.rs index 69df8e5e38..a3588e127a 100644 --- a/pyrefly/lib/lsp/non_wasm/workspace.rs +++ b/pyrefly/lib/lsp/non_wasm/workspace.rs @@ -238,6 +238,8 @@ pub struct DisabledLanguageServices { #[serde(default)] pub references: bool, #[serde(default)] + pub code_lens: bool, + #[serde(default)] pub rename: bool, #[serde(default)] pub signature_help: bool, @@ -248,8 +250,6 @@ pub struct DisabledLanguageServices { #[serde(default)] pub document_symbol: bool, #[serde(default)] - pub code_lens: bool, - #[serde(default)] pub semantic_tokens: bool, #[serde(default)] pub implementation: bool, @@ -267,12 +267,12 @@ impl DisabledLanguageServices { "textDocument/completion" => self.completion, "textDocument/documentHighlight" => self.document_highlight, "textDocument/references" => self.references, + "textDocument/codeLens" => self.code_lens, "textDocument/rename" => self.rename, "textDocument/signatureHelp" => self.signature_help, "textDocument/hover" => self.hover, "textDocument/inlayHint" => self.inlay_hint, "textDocument/documentSymbol" => self.document_symbol, - "textDocument/codeLens" => self.code_lens, "textDocument/semanticTokens/full" | "textDocument/semanticTokens/range" => { self.semantic_tokens } diff --git a/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs b/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs index 3a599239fd..08f2e5e096 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs @@ -6,8 +6,10 @@ */ use lsp_types::CodeLens; +use lsp_types::CodeLensOptions; use lsp_types::Url; use lsp_types::request::CodeLensRequest; +use pyrefly::commands::lsp::IndexingMode; use serde_json::Value; use serde_json::json; @@ -206,3 +208,85 @@ fn test_code_lens_disabled_by_default() { interaction.shutdown().unwrap(); } + +#[test] +fn test_initialize_advertises_code_lens_with_indexing() { + let interaction = LspInteraction::new_with_indexing_mode(IndexingMode::LazyBlocking); + + interaction + .client + .send_initialize( + interaction + .client + .get_initialize_params(&InitializeSettings::default()), + ) + .expect_response_with(|response| { + response.capabilities.code_lens_provider + == Some(CodeLensOptions { + resolve_provider: Some(false), + }) + }) + .unwrap(); + interaction.client.send_initialized(); + interaction.shutdown().unwrap(); +} + +#[test] +fn test_code_lens_shows_reference_counts() { + let root = get_test_files_root(); + let root_path = root.path().join("code_lens_references"); + let scope_uri = Url::from_file_path(&root_path).unwrap(); + let mut interaction = LspInteraction::new_with_indexing_mode(IndexingMode::LazyBlocking); + interaction.set_root(root_path); + interaction + .initialize(InitializeSettings { + workspace_folders: Some(vec![("test".to_owned(), scope_uri)]), + ..Default::default() + }) + .unwrap(); + + interaction.client.did_open("symbols.py"); + interaction.client.did_open("usage.py"); + + interaction + .client + .code_lens("symbols.py") + .expect_response_with(|response| { + let Some(lenses) = response else { + return false; + }; + has_reference_lens(&lenses, 0, "3 references", 3) + && has_reference_lens(&lenses, 1, "2 references", 2) + && has_reference_lens(&lenses, 4, "2 references", 2) + && lenses.len() == 3 + }) + .unwrap(); + + interaction.shutdown().unwrap(); +} + +fn has_reference_lens( + lenses: &[CodeLens], + line: u32, + expected_title: &str, + expected_locations: usize, +) -> bool { + lenses.iter().any(|lens| { + let Some(command) = &lens.command else { + return false; + }; + if lens.range.start.line != line + || command.title != expected_title + || command.command != "editor.action.showReferences" + { + return false; + } + let Some(arguments) = &command.arguments else { + return false; + }; + arguments + .get(2) + .and_then(|value| value.as_array()) + .is_some_and(|locations| locations.len() == expected_locations) + }) +} diff --git a/pyrefly/lib/test/lsp/lsp_interaction/object_model.rs b/pyrefly/lib/test/lsp/lsp_interaction/object_model.rs index 85e1417fc3..73ba5eea66 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/object_model.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/object_model.rs @@ -45,6 +45,7 @@ use lsp_types::notification::Initialized; use lsp_types::notification::Notification as _; use lsp_types::notification::PublishDiagnostics; use lsp_types::request::CodeActionRequest; +use lsp_types::request::CodeLensRequest; use lsp_types::request::Completion; use lsp_types::request::DocumentDiagnosticRequest; use lsp_types::request::DocumentHighlightRequest; @@ -519,6 +520,15 @@ impl TestClient { })) } + pub fn code_lens(&self, file: &'static str) -> ClientRequestHandle<'_, CodeLensRequest> { + let path = self.get_root_or_panic().join(file); + self.send_request(json!({ + "textDocument": { + "uri": Url::from_file_path(&path).unwrap().to_string() + } + })) + } + pub fn diagnostic( &self, file: &'static str, diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/code_lens_references/symbols.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/code_lens_references/symbols.py new file mode 100644 index 0000000000..7567cdd73d --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/code_lens_references/symbols.py @@ -0,0 +1,6 @@ +class Greeter: + def greet(self) -> str: + return "hi" + +def run(g: Greeter) -> str: + return g.greet() diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/code_lens_references/usage.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/code_lens_references/usage.py new file mode 100644 index 0000000000..924914d9ae --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/code_lens_references/usage.py @@ -0,0 +1,5 @@ +from symbols import Greeter, run + +g = Greeter() +g.greet() +run(g)