From 626f25f2517767776e67f528ab0aac65db6620ea Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Wed, 18 Mar 2026 10:23:42 +0900 Subject: [PATCH 1/3] fix --- pyrefly/lib/lsp/non_wasm/server.rs | 268 +++++++++++++++--- pyrefly/lib/lsp/non_wasm/workspace.rs | 6 +- .../lib/test/lsp/lsp_interaction/code_lens.rs | 84 ++++++ .../test/lsp/lsp_interaction/object_model.rs | 10 + .../code_lens_references/symbols.py | 6 + .../test_files/code_lens_references/usage.py | 5 + 6 files changed, 334 insertions(+), 45 deletions(-) create mode 100644 pyrefly/lib/test/lsp/lsp_interaction/test_files/code_lens_references/symbols.py create mode 100644 pyrefly/lib/test/lsp/lsp_interaction/test_files/code_lens_references/usage.py diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 6165ebf02a..20cd1375d0 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -43,6 +43,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; @@ -1054,6 +1055,12 @@ struct InitializeResult { server_info: Option, } +#[derive(Debug, Clone)] +struct CodeLensTarget { + range: Range, + definition: FindDefinitionItemWithDocstring, +} + pub fn initialize_finish( sender: &Sender, reader: &mut MessageReader, @@ -2228,18 +2235,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::( @@ -2286,6 +2281,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::( @@ -4679,6 +4691,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<'_>, @@ -4947,6 +5017,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_code_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, @@ -5111,36 +5325,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_code_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 dfc7e01e73..7db07362ec 100644 --- a/pyrefly/lib/lsp/non_wasm/workspace.rs +++ b/pyrefly/lib/lsp/non_wasm/workspace.rs @@ -284,6 +284,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, @@ -294,8 +296,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, @@ -313,12 +313,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 9c5809ffe6..9ccc27d7ab 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; @@ -520,6 +521,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) From 257c1b5a745fda4e8cc7872287f76a8b9b04503b Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 25 Jul 2026 11:37:49 +0900 Subject: [PATCH 2/3] comment --- crates/pyrefly_lsp_test/src/object_model.rs | 10 - pyrefly/lib/lsp/non_wasm.rs | 1 - pyrefly/lib/lsp/non_wasm/code_lens.rs | 205 -------- pyrefly/lib/lsp/non_wasm/server.rs | 231 +-------- pyrefly/lib/lsp/non_wasm/server/code_lens.rs | 475 ++++++++++++++++++ pyrefly/lib/lsp/non_wasm/workspace.rs | 6 +- .../lib/test/lsp/lsp_interaction/code_lens.rs | 42 +- 7 files changed, 495 insertions(+), 475 deletions(-) delete mode 100644 pyrefly/lib/lsp/non_wasm/code_lens.rs create mode 100644 pyrefly/lib/lsp/non_wasm/server/code_lens.rs diff --git a/crates/pyrefly_lsp_test/src/object_model.rs b/crates/pyrefly_lsp_test/src/object_model.rs index 9c59ae15cf..f74386c10f 100644 --- a/crates/pyrefly_lsp_test/src/object_model.rs +++ b/crates/pyrefly_lsp_test/src/object_model.rs @@ -56,7 +56,6 @@ use lsp_types::PublishDiagnosticsNotification; use lsp_types::PublishDiagnosticsParams; use lsp_types::ReferencesRequest; use lsp_types::RegistrationParams; -use lsp_types::CodeLensRequest; use lsp_types::RegistrationRequest; use lsp_types::RenameRequest; use lsp_types::Request as _; @@ -572,15 +571,6 @@ 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/lsp/non_wasm.rs b/pyrefly/lib/lsp/non_wasm.rs index 7b40188376..6a7c8dcbf0 100644 --- a/pyrefly/lib/lsp/non_wasm.rs +++ b/pyrefly/lib/lsp/non_wasm.rs @@ -7,7 +7,6 @@ mod build_system; pub mod call_hierarchy; -pub mod code_lens; pub mod connection; pub mod convert_module_package; pub mod document_symbols; diff --git a/pyrefly/lib/lsp/non_wasm/code_lens.rs b/pyrefly/lib/lsp/non_wasm/code_lens.rs deleted file mode 100644 index 594358e72b..0000000000 --- a/pyrefly/lib/lsp/non_wasm/code_lens.rs +++ /dev/null @@ -1,205 +0,0 @@ -/* - * 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. - */ - -use lsp_types::CodeLens; -use lsp_types::Command; -use lsp_types::Range; -use lsp_types::Uri; -use pyrefly_build::handle::Handle; -use pyrefly_python::ast::Ast; -use ruff_python_ast::Expr; -use ruff_python_ast::ExprAttribute; -use ruff_python_ast::Stmt; -use ruff_python_ast::StmtClassDef; -use ruff_text_size::TextRange; -use serde_json::Value; - -use crate::state::state::Transaction; - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum CodeLensKind { - Run, - Test, -} - -#[derive(Clone, Debug)] -pub struct CodeLensEntry { - pub range: TextRange, - pub kind: CodeLensKind, - pub test_name: Option, - pub class_name: Option, - pub is_unittest: bool, -} - -pub fn runnable_lsp_code_lens( - uri: &Uri, - range: Range, - entry: CodeLensEntry, - cwd: Option<&str>, -) -> CodeLens { - let (title, command, arguments) = match entry.kind { - CodeLensKind::Run => { - let mut args = serde_json::Map::new(); - args.insert("uri".to_owned(), serde_json::json!(uri.to_string())); - if let Some(cwd) = cwd { - args.insert("cwd".to_owned(), serde_json::json!(cwd)); - } - ("Run", "pyrefly.runMain", Some(vec![Value::Object(args)])) - } - CodeLensKind::Test => { - let mut args = serde_json::Map::new(); - args.insert("uri".to_owned(), serde_json::json!(uri.to_string())); - if let Some(cwd) = cwd { - args.insert("cwd".to_owned(), serde_json::json!(cwd)); - } - args.insert( - "position".to_owned(), - serde_json::json!({ - "line": range.start.line, - "character": range.start.character, - }), - ); - if let Some(test_name) = entry.test_name { - args.insert("testName".to_owned(), serde_json::json!(test_name)); - } - if let Some(class_name) = entry.class_name { - args.insert("className".to_owned(), serde_json::json!(class_name)); - } - args.insert( - "isUnittest".to_owned(), - serde_json::json!(entry.is_unittest), - ); - ("Test", "pyrefly.runTest", Some(vec![Value::Object(args)])) - } - }; - - CodeLens { - range, - command: Some(Command { - title: title.to_owned(), - command: command.to_owned(), - arguments, - tooltip: None, - }), - data: None, - } -} - -impl<'a> Transaction<'a> { - pub fn runnable_code_lens_entries( - &self, - handle: &Handle, - uri: &Uri, - runnable_code_lens: bool, - ) -> Option> { - if !runnable_code_lens || uri.path().ends_with(".pyi") || uri.path().ends_with(".ipynb") { - return Some(Vec::new()); - } - let ast = self.get_ast(handle)?; - let mut entries = Vec::new(); - collect_module_entries(&ast.body, &mut entries); - Some(entries) - } -} - -fn collect_module_entries(stmts: &[Stmt], entries: &mut Vec) { - for stmt in stmts { - match stmt { - Stmt::FunctionDef(func) => { - maybe_push_test(entries, func.name.as_str(), func.name.range, None, false); - } - Stmt::ClassDef(class_def) => { - let is_unittest = is_unittest_class(class_def); - if is_test_class(class_def, is_unittest) { - entries.push(CodeLensEntry { - range: class_def.name.range, - kind: CodeLensKind::Test, - test_name: None, - class_name: Some(class_def.name.as_str().to_owned()), - is_unittest, - }); - } - collect_class_entries( - &class_def.body, - entries, - class_def.name.as_str(), - is_unittest, - ); - } - Stmt::If(stmt_if) if Ast::is_main_guard(&stmt_if.test) => { - entries.push(CodeLensEntry { - range: stmt_if.range, - kind: CodeLensKind::Run, - test_name: None, - class_name: None, - is_unittest: false, - }); - } - _ => {} - } - } -} - -fn collect_class_entries( - stmts: &[Stmt], - entries: &mut Vec, - class_name: &str, - is_unittest: bool, -) { - for stmt in stmts { - if let Stmt::FunctionDef(func) = stmt { - maybe_push_test( - entries, - func.name.as_str(), - func.name.range, - Some(class_name.to_owned()), - is_unittest, - ); - } - } -} - -fn maybe_push_test( - entries: &mut Vec, - name: &str, - range: TextRange, - class_name: Option, - is_unittest: bool, -) { - if is_test_name(name) { - entries.push(CodeLensEntry { - range, - kind: CodeLensKind::Test, - test_name: Some(name.to_owned()), - class_name, - is_unittest, - }); - } -} - -fn is_test_name(name: &str) -> bool { - name.starts_with("test_") -} - -fn is_test_class(class_def: &StmtClassDef, is_unittest: bool) -> bool { - if class_def.name.as_str().starts_with("Test") { - return true; - } - is_unittest -} - -fn is_unittest_class(class_def: &StmtClassDef) -> bool { - class_def.bases().iter().any(is_unittest_base) -} - -fn is_unittest_base(base: &Expr) -> bool { - match base { - Expr::Name(name) => name.id.as_str().ends_with("TestCase"), - Expr::Attribute(ExprAttribute { attr, .. }) => attr.id.as_str().ends_with("TestCase"), - _ => false, - } -} diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index daf331f82b..cfafcf72f2 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -6,6 +6,7 @@ */ mod event_loop; +mod code_lens; use std::cmp::min; use std::collections::HashMap; @@ -52,11 +53,8 @@ use lsp_types::CodeActionProvider; use lsp_types::CodeActionRequest; use lsp_types::CodeActionResponse; use lsp_types::CodeActionTriggerKind; -use lsp_types::CodeLens; use lsp_types::CodeLensOptions; -use lsp_types::CodeLensParams; use lsp_types::CodeLensRequest; -use lsp_types::Command; use lsp_types::CompletionItem; use lsp_types::CompletionList; use lsp_types::CompletionOptions; @@ -300,7 +298,6 @@ 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::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; use crate::lsp::non_wasm::external_provider::ExternalProvider; @@ -1420,12 +1417,6 @@ struct InitializeResult { server_info: Option, } -#[derive(Debug, Clone)] -struct CodeLensTarget { - range: Range, - definition: FindDefinitionItemWithDocstring, -} - pub fn initialize_finish( sender: &Sender, reader: &mut MessageReader, @@ -3202,24 +3193,6 @@ impl Server { telemetry.set_file_stats(TelemetryFileStats { uri, config_root }); } - fn runnable_code_lens_cwd(&self, path: &std::path::Path) -> Option { - let config = self.state.config_finder().python_file( - ModuleNameWithKind::guaranteed(ModuleName::unknown()), - &ModulePath::filesystem(path.to_path_buf()), - ); - let cwd = config - .source - .root_from_file() - .map(std::path::Path::to_path_buf) - .or_else(|| { - self.workspaces - .get_with(path.to_path_buf(), |(workspace_root, _)| { - workspace_root.cloned() - }) - })?; - Some(cwd.to_string_lossy().into_owned()) - } - fn send_response(&self, x: Response) { self.connection.send(Message::Response(x)) } @@ -5530,64 +5503,6 @@ 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<'_>, @@ -5863,150 +5778,6 @@ 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_code_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, diff --git a/pyrefly/lib/lsp/non_wasm/server/code_lens.rs b/pyrefly/lib/lsp/non_wasm/server/code_lens.rs new file mode 100644 index 0000000000..f2424fc44e --- /dev/null +++ b/pyrefly/lib/lsp/non_wasm/server/code_lens.rs @@ -0,0 +1,475 @@ +/* + * 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. + */ + +use dupe::Dupe; +use lsp_server::ErrorCode; +use lsp_server::RequestId; +use lsp_types::CodeLens; +use lsp_types::CodeLensParams; +use lsp_types::Command; +use lsp_types::DocumentSymbol; +use lsp_types::Location; +use lsp_types::Range; +use lsp_types::SymbolKind; +use lsp_types::Url; +use lsp_types::request::CodeLensRequest; +use lsp_types::request::Request as _; +use pyrefly_build::handle::Handle; +use pyrefly_python::ast::Ast; +use pyrefly_python::module::TextRangeWithModule; +use pyrefly_python::module_name::ModuleName; +use pyrefly_python::module_name::ModuleNameWithKind; +use pyrefly_python::module_path::ModulePath; +use pyrefly_util::task_heap::Cancelled; +use pyrefly_util::telemetry::ActivityKey; +use pyrefly_util::telemetry::EmptyResponseReason; +use pyrefly_util::telemetry::TelemetryEventKind; +use ruff_python_ast::Expr; +use ruff_python_ast::ExprAttribute; +use ruff_python_ast::Stmt; +use ruff_python_ast::StmtClassDef; +use ruff_text_size::TextRange; +use serde_json::Value; +use starlark_map::small_set::SmallSet; +use tracing::info; + +use super::Server; +use crate::commands::lsp::IndexingMode; +use crate::lsp::non_wasm::lsp::new_response; +use crate::lsp::non_wasm::module_helpers::PathRemapper; +use crate::lsp::non_wasm::module_helpers::module_info_to_uri; +use crate::lsp::non_wasm::protocol::Message; +use crate::lsp::non_wasm::protocol::Response; +use crate::state::lsp::FindDefinitionItemWithDocstring; +use crate::state::lsp::FindPreference; +use crate::state::lsp::ImportBehavior; +use crate::state::state::CancellableTransaction; +use crate::state::state::Transaction; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum CodeLensKind { + Run, + Test, +} + +#[derive(Clone, Debug)] +struct CodeLensEntry { + range: TextRange, + kind: CodeLensKind, + test_name: Option, + class_name: Option, + is_unittest: bool, +} + +#[derive(Debug)] +struct CodeLensTarget { + range: Range, + definition: FindDefinitionItemWithDocstring, +} + +impl Server { + /// Build runnable lenses and resolve reference counts on the references queue. + pub(super) 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.reference_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 reference_lenses = match resolve_reference_code_lenses( + &mut transaction, + &handle, + &source_uri, + targets, + path_remapper.as_ref(), + ) { + Ok(lenses) => lenses, + 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 lenses = runnable_lenses; + lenses.extend(reference_lenses); + server.cancellation_handles.lock().remove(&request_id); + server.connection.send(Message::Response(new_response( + request_id, + Ok(Some(lenses)), + ))); + }), + ); + Ok(()) + } + + fn reference_code_lens_targets( + &self, + transaction: &Transaction<'_>, + handle: &Handle, + uri: &Url, + ) -> Result, EmptyResponseReason> { + 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(); + collect_reference_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) { + targets.push(CodeLensTarget { + range: symbol.selection_range, + definition, + }); + } + } + Ok(targets) + } + + 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_code_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 config = self.state.config_finder().python_file( + ModuleNameWithKind::guaranteed(ModuleName::unknown()), + &ModulePath::filesystem(path.clone()), + ); + let cwd = config + .source + .root_from_file() + .map(std::path::Path::to_path_buf) + .or_else(|| { + self.workspaces + .get_with(path, |(workspace_root, _)| workspace_root.cloned()) + }) + .map(|path| path.to_string_lossy().into_owned()); + + 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 collect_reference_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() { + collect_reference_symbols(children, out); + } + } +} + +fn resolve_reference_code_lenses( + transaction: &mut CancellableTransaction<'_>, + handle: &Handle, + source_uri: &Url, + targets: Vec, + path_remapper: Option<&PathRemapper>, +) -> Result, Cancelled> { + let mut lenses = Vec::with_capacity(targets.len()); + for target in targets { + let local_results = transaction.find_global_references_from_definition( + *handle.sys_info(), + target.definition.metadata, + TextRangeWithModule::new( + target.definition.module.clone(), + target.definition.definition_range, + ), + false, + )?; + + let mut locations = Vec::new(); + for (info, ranges) in local_results { + if let Some(uri) = module_info_to_uri(&info, path_remapper) { + 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"), + ]), + tooltip: None, + }), + data: None, + }); + } + Ok(lenses) +} + +fn runnable_lsp_code_lens( + uri: &Url, + range: Range, + entry: CodeLensEntry, + cwd: Option<&str>, +) -> CodeLens { + let (title, command, arguments) = match entry.kind { + CodeLensKind::Run => { + let mut args = serde_json::Map::new(); + args.insert("uri".to_owned(), serde_json::json!(uri.to_string())); + if let Some(cwd) = cwd { + args.insert("cwd".to_owned(), serde_json::json!(cwd)); + } + ("Run", "pyrefly.runMain", Some(vec![Value::Object(args)])) + } + CodeLensKind::Test => { + let mut args = serde_json::Map::new(); + args.insert("uri".to_owned(), serde_json::json!(uri.to_string())); + if let Some(cwd) = cwd { + args.insert("cwd".to_owned(), serde_json::json!(cwd)); + } + args.insert( + "position".to_owned(), + serde_json::json!({ + "line": range.start.line, + "character": range.start.character, + }), + ); + if let Some(test_name) = entry.test_name { + args.insert("testName".to_owned(), serde_json::json!(test_name)); + } + if let Some(class_name) = entry.class_name { + args.insert("className".to_owned(), serde_json::json!(class_name)); + } + args.insert( + "isUnittest".to_owned(), + serde_json::json!(entry.is_unittest), + ); + ("Test", "pyrefly.runTest", Some(vec![Value::Object(args)])) + } + }; + + CodeLens { + range, + command: Some(Command { + title: title.to_owned(), + command: command.to_owned(), + arguments, + tooltip: None, + }), + data: None, + } +} + +impl<'a> Transaction<'a> { + fn runnable_code_lens_entries( + &self, + handle: &Handle, + uri: &Url, + runnable_code_lens: bool, + ) -> Option> { + if !runnable_code_lens || uri.path().ends_with(".pyi") || uri.path().ends_with(".ipynb") { + return Some(Vec::new()); + } + let ast = self.get_ast(handle)?; + let mut entries = Vec::new(); + collect_module_entries(&ast.body, &mut entries); + Some(entries) + } +} + +fn collect_module_entries(stmts: &[Stmt], entries: &mut Vec) { + for stmt in stmts { + match stmt { + Stmt::FunctionDef(func) => { + maybe_push_test(entries, func.name.as_str(), func.name.range, None, false); + } + Stmt::ClassDef(class_def) => { + let is_unittest = is_unittest_class(class_def); + if is_test_class(class_def, is_unittest) { + entries.push(CodeLensEntry { + range: class_def.name.range, + kind: CodeLensKind::Test, + test_name: None, + class_name: Some(class_def.name.as_str().to_owned()), + is_unittest, + }); + } + collect_class_entries( + &class_def.body, + entries, + class_def.name.as_str(), + is_unittest, + ); + } + Stmt::If(stmt_if) if Ast::is_main_guard(&stmt_if.test) => { + entries.push(CodeLensEntry { + range: stmt_if.range, + kind: CodeLensKind::Run, + test_name: None, + class_name: None, + is_unittest: false, + }); + } + _ => {} + } + } +} + +fn collect_class_entries( + stmts: &[Stmt], + entries: &mut Vec, + class_name: &str, + is_unittest: bool, +) { + for stmt in stmts { + if let Stmt::FunctionDef(func) = stmt { + maybe_push_test( + entries, + func.name.as_str(), + func.name.range, + Some(class_name.to_owned()), + is_unittest, + ); + } + } +} + +fn maybe_push_test( + entries: &mut Vec, + name: &str, + range: TextRange, + class_name: Option, + is_unittest: bool, +) { + if is_test_name(name) { + entries.push(CodeLensEntry { + range, + kind: CodeLensKind::Test, + test_name: Some(name.to_owned()), + class_name, + is_unittest, + }); + } +} + +fn is_test_name(name: &str) -> bool { + name.starts_with("test_") +} + +fn is_test_class(class_def: &StmtClassDef, is_unittest: bool) -> bool { + if class_def.name.as_str().starts_with("Test") { + return true; + } + is_unittest +} + +fn is_unittest_class(class_def: &StmtClassDef) -> bool { + class_def.bases().iter().any(is_unittest_base) +} + +fn is_unittest_base(base: &Expr) -> bool { + match base { + Expr::Name(name) => name.id.as_str().ends_with("TestCase"), + Expr::Attribute(ExprAttribute { attr, .. }) => attr.id.as_str().ends_with("TestCase"), + _ => false, + } +} diff --git a/pyrefly/lib/lsp/non_wasm/workspace.rs b/pyrefly/lib/lsp/non_wasm/workspace.rs index 9f0ec6f7e5..1e109d34e5 100644 --- a/pyrefly/lib/lsp/non_wasm/workspace.rs +++ b/pyrefly/lib/lsp/non_wasm/workspace.rs @@ -295,8 +295,6 @@ 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, @@ -307,6 +305,8 @@ 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, @@ -324,12 +324,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 6b38b2aa60..2ca23880c9 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs @@ -6,12 +6,13 @@ */ use lsp_types::CodeLens; -use lsp_types::CodeLensOptions; use lsp_types::CodeLensRequest; use lsp_types::Uri; +use pyrefly_lsp_test::IndexingMode; +use pyrefly_lsp_test::LspArgs; use pyrefly_lsp_test::object_model::InitializeSettings; use pyrefly_lsp_test::object_model::LspInteraction; -use pyrefly_lsp_test::IndexingMode; +use pyrefly_lsp_test::object_model::LspInteractionArgs; use serde_json::Value; use serde_json::json; @@ -209,34 +210,19 @@ 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); + let symbols_uri = Url::from_file_path(root_path.join("symbols.py")).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); interaction .initialize(InitializeSettings { @@ -250,7 +236,11 @@ fn test_code_lens_shows_reference_counts() { interaction .client - .code_lens("symbols.py") + .send_request::(json!({ + "textDocument": { + "uri": symbols_uri + } + })) .expect_response_with(|response| { let Some(lenses) = response else { return false; From 0feac379e3839ec0283e681a8b766ab40201c7cb Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 28 Jul 2026 11:45:48 +0900 Subject: [PATCH 3/3] comment --- pyrefly/lib/lsp/non_wasm/server.rs | 32 +- pyrefly/lib/lsp/non_wasm/server/code_lens.rs | 516 +++--------------- .../server/code_lens/reference_count.rs | 121 ++++ .../lsp/non_wasm/server/code_lens/runnable.rs | 254 +++++++++ pyrefly/lib/test/lsp/lsp_interaction/basic.rs | 2 +- .../lib/test/lsp/lsp_interaction/code_lens.rs | 218 ++++++-- 6 files changed, 650 insertions(+), 493 deletions(-) create mode 100644 pyrefly/lib/lsp/non_wasm/server/code_lens/reference_count.rs create mode 100644 pyrefly/lib/lsp/non_wasm/server/code_lens/runnable.rs diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index cfafcf72f2..23e08a2da2 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -5,8 +5,8 @@ * LICENSE file in the root directory of this source tree. */ -mod event_loop; mod code_lens; +mod event_loop; use std::cmp::min; use std::collections::HashMap; @@ -55,6 +55,7 @@ use lsp_types::CodeActionResponse; use lsp_types::CodeActionTriggerKind; use lsp_types::CodeLensOptions; use lsp_types::CodeLensRequest; +use lsp_types::CodeLensResolveRequest; use lsp_types::CompletionItem; use lsp_types::CompletionList; use lsp_types::CompletionOptions; @@ -1670,7 +1671,7 @@ pub fn capabilities( ..Default::default() })), code_lens_provider: Some(CodeLensOptions { - resolve_provider: Some(false), + resolve_provider: Some(true), work_done_progress_options: Default::default(), }), completion_provider: Some(CompletionOptions { @@ -2374,6 +2375,7 @@ impl Server { DefinitionRequest::METHOD.as_str(), ProvideType::METHOD.as_str(), TypeErrorDisplayStatusRequest::METHOD.as_str(), + CodeLensResolveRequest::METHOD.as_str(), ]; let in_cancelled_requests = canceled_requests.remove(&x.id); @@ -2750,15 +2752,29 @@ impl Server { ) { self.set_file_stats(params.text_document.uri.clone(), telemetry_event); - if let Err(reason) = self.code_lens( + let response = match self.code_lens(&transaction, params) { + Ok(response) => response, + Err(reason) => { + telemetry_event.set_empty_response_reason(reason); + None + } + }; + 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, + ) + && let Err(reason) = self.resolve_code_lens( x.id.clone(), &transaction, - params, + ¶ms, telemetry_event.activity_key.clone(), - ) { - self.send_response(new_response(x.id, Ok(None::<()>))); - telemetry_event.set_empty_response_reason(reason); - } + ) + { + telemetry_event.set_empty_response_reason(reason); + self.send_response(new_response(x.id, Ok(params))); } } else if let Some(params) = as_request::(&x) { if let Some(params) = self diff --git a/pyrefly/lib/lsp/non_wasm/server/code_lens.rs b/pyrefly/lib/lsp/non_wasm/server/code_lens.rs index f2424fc44e..4c383cf376 100644 --- a/pyrefly/lib/lsp/non_wasm/server/code_lens.rs +++ b/pyrefly/lib/lsp/non_wasm/server/code_lens.rs @@ -5,471 +5,117 @@ * LICENSE file in the root directory of this source tree. */ -use dupe::Dupe; +//! Server orchestration for runnable and reference-count CodeLens providers. + +mod reference_count; +mod runnable; + use lsp_server::ErrorCode; use lsp_server::RequestId; use lsp_types::CodeLens; use lsp_types::CodeLensParams; -use lsp_types::Command; -use lsp_types::DocumentSymbol; -use lsp_types::Location; -use lsp_types::Range; -use lsp_types::SymbolKind; -use lsp_types::Url; -use lsp_types::request::CodeLensRequest; -use lsp_types::request::Request as _; -use pyrefly_build::handle::Handle; -use pyrefly_python::ast::Ast; +use lsp_types::CodeLensRequest; +use lsp_types::Request as _; use pyrefly_python::module::TextRangeWithModule; -use pyrefly_python::module_name::ModuleName; -use pyrefly_python::module_name::ModuleNameWithKind; -use pyrefly_python::module_path::ModulePath; -use pyrefly_util::task_heap::Cancelled; use pyrefly_util::telemetry::ActivityKey; use pyrefly_util::telemetry::EmptyResponseReason; -use pyrefly_util::telemetry::TelemetryEventKind; -use ruff_python_ast::Expr; -use ruff_python_ast::ExprAttribute; -use ruff_python_ast::Stmt; -use ruff_python_ast::StmtClassDef; -use ruff_text_size::TextRange; -use serde_json::Value; -use starlark_map::small_set::SmallSet; -use tracing::info; use super::Server; use crate::commands::lsp::IndexingMode; -use crate::lsp::non_wasm::lsp::new_response; -use crate::lsp::non_wasm::module_helpers::PathRemapper; -use crate::lsp::non_wasm::module_helpers::module_info_to_uri; use crate::lsp::non_wasm::protocol::Message; use crate::lsp::non_wasm::protocol::Response; use crate::state::lsp::FindDefinitionItemWithDocstring; use crate::state::lsp::FindPreference; use crate::state::lsp::ImportBehavior; -use crate::state::state::CancellableTransaction; +use crate::state::lsp::ReferenceOptions; use crate::state::state::Transaction; -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum CodeLensKind { - Run, - Test, -} - -#[derive(Clone, Debug)] -struct CodeLensEntry { - range: TextRange, - kind: CodeLensKind, - test_name: Option, - class_name: Option, - is_unittest: bool, -} - -#[derive(Debug)] -struct CodeLensTarget { - range: Range, - definition: FindDefinitionItemWithDocstring, -} - impl Server { - /// Build runnable lenses and resolve reference counts on the references queue. - pub(super) fn code_lens<'a>( - &'a self, - request_id: RequestId, - transaction: &Transaction<'a>, + /// Return eager runnable lenses and unresolved reference-count lenses for a Python file. + pub(super) fn code_lens( + &self, + transaction: &Transaction<'_>, params: CodeLensParams, - activity_key: Option, - ) -> Result<(), EmptyResponseReason> { + ) -> 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(()); + return Ok(Some(Vec::new())); } - 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.reference_code_lens_targets(transaction, &handle, uri)? - }; - if targets.is_empty() { - self.send_response(new_response(request_id, Ok(Some(runnable_lenses)))); - return Ok(()); + let handle = self.make_handle_if_enabled(uri, Some(CodeLensRequest::METHOD.as_str()))?; + let mut lenses = runnable::code_lenses(self, transaction, &handle, uri)?; + if self.indexing_mode != IndexingMode::None { + lenses.extend(reference_count::unresolved_code_lenses( + transaction, + &handle, + uri, + )?); } + Ok(Some(lenses)) + } - 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 reference_lenses = match resolve_reference_code_lenses( - &mut transaction, - &handle, - &source_uri, - targets, - path_remapper.as_ref(), - ) { - Ok(lenses) => lenses, - 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 lenses = runnable_lenses; - lenses.extend(reference_lenses); - server.cancellation_handles.lock().remove(&request_id); - server.connection.send(Message::Response(new_response( + /// Resolve one visible reference-count lens on the references queue. + pub(super) fn resolve_code_lens<'a>( + &'a self, + request_id: RequestId, + transaction: &Transaction<'a>, + lens: &CodeLens, + activity_key: Option, + ) -> Result<(), EmptyResponseReason> { + let data = match reference_count::resolve_data(lens) { + Ok(data) => data, + Err(message) => { + self.connection.send(Message::Response(Response::new_err( request_id, - Ok(Some(lenses)), + ErrorCode::InvalidParams as i32, + message, ))); - }), - ); - Ok(()) - } - - fn reference_code_lens_targets( - &self, - transaction: &Transaction<'_>, - handle: &Handle, - uri: &Url, - ) -> Result, EmptyResponseReason> { - 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(); - collect_reference_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) { - targets.push(CodeLensTarget { - range: symbol.selection_range, - definition, - }); + return Ok(()); } - } - Ok(targets) - } - - 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_code_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 config = self.state.config_finder().python_file( - ModuleNameWithKind::guaranteed(ModuleName::unknown()), - &ModulePath::filesystem(path.clone()), - ); - let cwd = config - .source - .root_from_file() - .map(std::path::Path::to_path_buf) - .or_else(|| { - self.workspaces - .get_with(path, |(workspace_root, _)| workspace_root.cloned()) - }) - .map(|path| path.to_string_lossy().into_owned()); - - 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 collect_reference_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() { - collect_reference_symbols(children, out); - } - } -} - -fn resolve_reference_code_lenses( - transaction: &mut CancellableTransaction<'_>, - handle: &Handle, - source_uri: &Url, - targets: Vec, - path_remapper: Option<&PathRemapper>, -) -> Result, Cancelled> { - let mut lenses = Vec::with_capacity(targets.len()); - for target in targets { - let local_results = transaction.find_global_references_from_definition( - *handle.sys_info(), - target.definition.metadata, - TextRangeWithModule::new( - target.definition.module.clone(), - target.definition.definition_range, - ), - false, - )?; - - let mut locations = Vec::new(); - for (info, ranges) in local_results { - if let Some(uri) = module_info_to_uri(&info, path_remapper) { - 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"), - ]), - tooltip: None, - }), - data: None, - }); - } - Ok(lenses) -} - -fn runnable_lsp_code_lens( - uri: &Url, - range: Range, - entry: CodeLensEntry, - cwd: Option<&str>, -) -> CodeLens { - let (title, command, arguments) = match entry.kind { - CodeLensKind::Run => { - let mut args = serde_json::Map::new(); - args.insert("uri".to_owned(), serde_json::json!(uri.to_string())); - if let Some(cwd) = cwd { - args.insert("cwd".to_owned(), serde_json::json!(cwd)); - } - ("Run", "pyrefly.runMain", Some(vec![Value::Object(args)])) + if self.open_notebook_cells.read().contains_key(&data.uri) { + return Err(EmptyResponseReason::NotebookNotSupported); } - CodeLensKind::Test => { - let mut args = serde_json::Map::new(); - args.insert("uri".to_owned(), serde_json::json!(uri.to_string())); - if let Some(cwd) = cwd { - args.insert("cwd".to_owned(), serde_json::json!(cwd)); - } - args.insert( - "position".to_owned(), - serde_json::json!({ - "line": range.start.line, - "character": range.start.character, - }), - ); - if let Some(test_name) = entry.test_name { - args.insert("testName".to_owned(), serde_json::json!(test_name)); - } - if let Some(class_name) = entry.class_name { - args.insert("className".to_owned(), serde_json::json!(class_name)); - } - args.insert( - "isUnittest".to_owned(), - serde_json::json!(entry.is_unittest), - ); - ("Test", "pyrefly.runTest", Some(vec![Value::Object(args)])) - } - }; - - CodeLens { - range, - command: Some(Command { - title: title.to_owned(), - command: command.to_owned(), - arguments, - tooltip: None, - }), - data: None, - } -} -impl<'a> Transaction<'a> { - fn runnable_code_lens_entries( - &self, - handle: &Handle, - uri: &Url, - runnable_code_lens: bool, - ) -> Option> { - if !runnable_code_lens || uri.path().ends_with(".pyi") || uri.path().ends_with(".ipynb") { - return Some(Vec::new()); - } - let ast = self.get_ast(handle)?; - let mut entries = Vec::new(); - collect_module_entries(&ast.body, &mut entries); - Some(entries) - } -} - -fn collect_module_entries(stmts: &[Stmt], entries: &mut Vec) { - for stmt in stmts { - match stmt { - Stmt::FunctionDef(func) => { - maybe_push_test(entries, func.name.as_str(), func.name.range, None, false); - } - Stmt::ClassDef(class_def) => { - let is_unittest = is_unittest_class(class_def); - if is_test_class(class_def, is_unittest) { - entries.push(CodeLensEntry { - range: class_def.name.range, - kind: CodeLensKind::Test, - test_name: None, - class_name: Some(class_def.name.as_str().to_owned()), - is_unittest, - }); - } - collect_class_entries( - &class_def.body, - entries, - class_def.name.as_str(), - is_unittest, - ); - } - Stmt::If(stmt_if) if Ast::is_main_guard(&stmt_if.test) => { - entries.push(CodeLensEntry { - range: stmt_if.range, - kind: CodeLensKind::Run, - test_name: None, - class_name: None, - is_unittest: false, - }); - } - _ => {} - } - } -} - -fn collect_class_entries( - stmts: &[Stmt], - entries: &mut Vec, - class_name: &str, - is_unittest: bool, -) { - for stmt in stmts { - if let Stmt::FunctionDef(func) = stmt { - maybe_push_test( - entries, - func.name.as_str(), - func.name.range, - Some(class_name.to_owned()), - is_unittest, - ); - } - } -} - -fn maybe_push_test( - entries: &mut Vec, - name: &str, - range: TextRange, - class_name: Option, - is_unittest: bool, -) { - if is_test_name(name) { - entries.push(CodeLensEntry { - range, - kind: CodeLensKind::Test, - test_name: Some(name.to_owned()), - class_name, - is_unittest, - }); - } -} - -fn is_test_name(name: &str) -> bool { - name.starts_with("test_") -} - -fn is_test_class(class_def: &StmtClassDef, is_unittest: bool) -> bool { - if class_def.name.as_str().starts_with("Test") { - return true; - } - is_unittest -} - -fn is_unittest_class(class_def: &StmtClassDef) -> bool { - class_def.bases().iter().any(is_unittest_base) -} - -fn is_unittest_base(base: &Expr) -> bool { - match base { - Expr::Name(name) => name.id.as_str().ends_with("TestCase"), - Expr::Attribute(ExprAttribute { attr, .. }) => attr.id.as_str().ends_with("TestCase"), - _ => false, + let handle = + self.make_handle_if_enabled(&data.uri, Some(CodeLensRequest::METHOD.as_str()))?; + let path_remapper = self.path_remapper.clone(); + let source_uri = data.uri; + let source_uri_for_response = source_uri.clone(); + let lens = lens.clone(); + self.async_find_from_definition_helper( + request_id, + transaction, + handle, + &source_uri, + lens.range.start, + FindPreference { + import_behavior: ImportBehavior::StopAtRenamedImports, + ..Default::default() + }, + activity_key, + move |transaction, handle, definition, _telemetry, _telemetry_event| { + let FindDefinitionItemWithDocstring { + metadata, + definition_range, + module, + .. + } = definition; + Ok(transaction.find_global_references_from_definition( + *handle.sys_info(), + metadata, + TextRangeWithModule::new(module, definition_range), + ReferenceOptions::all(false), + )?) + }, + move |results| { + reference_count::resolve_code_lens( + lens, + &source_uri_for_response, + results, + path_remapper.as_ref(), + ) + }, + ) } } diff --git a/pyrefly/lib/lsp/non_wasm/server/code_lens/reference_count.rs b/pyrefly/lib/lsp/non_wasm/server/code_lens/reference_count.rs new file mode 100644 index 0000000000..bddba10273 --- /dev/null +++ b/pyrefly/lib/lsp/non_wasm/server/code_lens/reference_count.rs @@ -0,0 +1,121 @@ +/* + * 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. + */ + +//! Lazy reference-count CodeLens creation and resolution. + +use lsp_types::CodeLens; +use lsp_types::Command; +use lsp_types::DocumentSymbol; +use lsp_types::Location; +use lsp_types::SymbolKind; +use lsp_types::Uri; +use pyrefly_build::handle::Handle; +use pyrefly_util::telemetry::EmptyResponseReason; +use ruff_text_size::TextRange; +use serde::Deserialize; +use serde::Serialize; + +use crate::ModuleInfo; +use crate::lsp::non_wasm::module_helpers::PathRemapper; +use crate::lsp::non_wasm::module_helpers::module_info_to_uri; +use crate::state::state::Transaction; + +#[derive(Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub(super) struct ReferenceCodeLensData { + pub(super) uri: Uri, +} + +/// Collect unresolved lenses without performing definition lookup or global reference searches. +pub(super) fn unresolved_code_lenses( + transaction: &Transaction<'_>, + handle: &Handle, + uri: &Uri, +) -> Result, EmptyResponseReason> { + let symbols = transaction + .symbols(handle, None) + .ok_or(EmptyResponseReason::ModuleInfoNotFound)?; + let mut lenses = Vec::new(); + collect_symbol_lenses(&symbols, uri, &mut lenses); + Ok(lenses) +} + +/// Decode the opaque data attached to an unresolved reference-count lens. +pub(super) fn resolve_data(lens: &CodeLens) -> Result { + let data = lens + .data + .clone() + .ok_or_else(|| "Reference CodeLens is missing resolve data".to_owned())?; + serde_json::from_value(data) + .map_err(|error| format!("Invalid reference CodeLens resolve data: {error}")) +} + +/// Attach the reference command, excluding every location backed by a notebook cell. +pub(super) fn resolve_code_lens( + mut lens: CodeLens, + source_uri: &Uri, + results: Vec<(ModuleInfo, Vec)>, + path_remapper: Option<&PathRemapper>, +) -> CodeLens { + let mut locations = Vec::new(); + for (info, ranges) in results { + let Some(uri) = module_info_to_uri(&info, path_remapper) else { + continue; + }; + for range in ranges { + if info.to_cell_for_lsp(range.start()).is_some() { + continue; + } + 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") + }; + lens.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(lens.range.start) + .expect("Position should serialize for code lens"), + serde_json::to_value(&locations).expect("Locations should serialize for code lens"), + ]), + tooltip: None, + }); + lens.data = None; + lens +} + +/// Traverse document symbols and create one unresolved lens per reference-countable definition. +fn collect_symbol_lenses(symbols: &[DocumentSymbol], uri: &Uri, lenses: &mut Vec) { + for symbol in symbols { + if matches!( + symbol.kind, + SymbolKind::Class | SymbolKind::Function | SymbolKind::Method + ) { + lenses.push(CodeLens { + range: symbol.selection_range, + command: None, + data: Some( + serde_json::to_value(ReferenceCodeLensData { uri: uri.clone() }) + .expect("Reference CodeLens data should serialize"), + ), + }); + } + if let Some(children) = symbol.children.as_deref() { + collect_symbol_lenses(children, uri, lenses); + } + } +} diff --git a/pyrefly/lib/lsp/non_wasm/server/code_lens/runnable.rs b/pyrefly/lib/lsp/non_wasm/server/code_lens/runnable.rs new file mode 100644 index 0000000000..6084caa3ed --- /dev/null +++ b/pyrefly/lib/lsp/non_wasm/server/code_lens/runnable.rs @@ -0,0 +1,254 @@ +/* + * 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. + */ + +//! Detection and rendering of Run and Test CodeLens commands. + +use lsp_types::CodeLens; +use lsp_types::Command; +use lsp_types::Range; +use lsp_types::Uri; +use pyrefly_build::handle::Handle; +use pyrefly_python::ast::Ast; +use pyrefly_python::module_name::ModuleName; +use pyrefly_python::module_name::ModuleNameWithKind; +use pyrefly_python::module_path::ModulePath; +use pyrefly_util::telemetry::EmptyResponseReason; +use ruff_python_ast::Expr; +use ruff_python_ast::ExprAttribute; +use ruff_python_ast::Stmt; +use ruff_python_ast::StmtClassDef; +use ruff_text_size::TextRange; +use serde_json::Value; + +use super::super::Server; +use crate::state::state::Transaction; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RunnableKind { + Run, + Test, +} + +#[derive(Clone, Debug)] +struct RunnableEntry { + range: TextRange, + kind: RunnableKind, + test_name: Option, + class_name: Option, + is_unittest: bool, +} + +/// Collect and render all enabled runnable lenses for one Python file. +pub(super) fn code_lenses( + server: &Server, + transaction: &Transaction<'_>, + handle: &Handle, + uri: &Uri, +) -> Result, EmptyResponseReason> { + let path = server + .path_for_uri(uri) + .ok_or(EmptyResponseReason::NoFilePath)?; + let runnable_code_lens = server + .workspaces + .get_with(path.clone(), |(_, workspace)| workspace.runnable_code_lens); + let maybe_cell_idx = server.maybe_get_code_cell_index(uri); + let info = transaction + .get_module_info(handle) + .ok_or(EmptyResponseReason::ModuleInfoNotFound)?; + let entries = runnable_entries(transaction, handle, uri, runnable_code_lens) + .ok_or(EmptyResponseReason::ModuleInfoNotFound)?; + let config = server.state.config_finder().python_file( + ModuleNameWithKind::guaranteed(ModuleName::unknown()), + &ModulePath::filesystem(path.clone()), + ); + let cwd = config + .source + .root_from_file() + .map(std::path::Path::to_path_buf) + .or_else(|| { + server + .workspaces + .get_with(path, |(workspace_root, _)| workspace_root.cloned()) + }) + .map(|path| path.to_string_lossy().into_owned()); + + 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(entry_to_lsp_code_lens(uri, range, entry, cwd.as_deref())); + } + Ok(lenses) +} + +/// Render a semantic runnable entry as an executable LSP command. +fn entry_to_lsp_code_lens( + uri: &Uri, + range: Range, + entry: RunnableEntry, + cwd: Option<&str>, +) -> CodeLens { + let (title, command, arguments) = match entry.kind { + RunnableKind::Run => { + let mut args = serde_json::Map::new(); + args.insert("uri".to_owned(), serde_json::json!(uri.to_string())); + if let Some(cwd) = cwd { + args.insert("cwd".to_owned(), serde_json::json!(cwd)); + } + ("Run", "pyrefly.runMain", Some(vec![Value::Object(args)])) + } + RunnableKind::Test => { + let mut args = serde_json::Map::new(); + args.insert("uri".to_owned(), serde_json::json!(uri.to_string())); + if let Some(cwd) = cwd { + args.insert("cwd".to_owned(), serde_json::json!(cwd)); + } + args.insert( + "position".to_owned(), + serde_json::json!({ + "line": range.start.line, + "character": range.start.character, + }), + ); + if let Some(test_name) = entry.test_name { + args.insert("testName".to_owned(), serde_json::json!(test_name)); + } + if let Some(class_name) = entry.class_name { + args.insert("className".to_owned(), serde_json::json!(class_name)); + } + args.insert( + "isUnittest".to_owned(), + serde_json::json!(entry.is_unittest), + ); + ("Test", "pyrefly.runTest", Some(vec![Value::Object(args)])) + } + }; + + CodeLens { + range, + command: Some(Command { + title: title.to_owned(), + command: command.to_owned(), + arguments, + tooltip: None, + }), + data: None, + } +} + +/// Extract runnable entries from the file AST without constructing LSP values. +fn runnable_entries( + transaction: &Transaction<'_>, + handle: &Handle, + uri: &Uri, + runnable_code_lens: bool, +) -> Option> { + if !runnable_code_lens || uri.path().ends_with(".pyi") || uri.path().ends_with(".ipynb") { + return Some(Vec::new()); + } + let ast = transaction.get_ast(handle)?; + let mut entries = Vec::new(); + collect_module_entries(&ast.body, &mut entries); + Some(entries) +} + +fn collect_module_entries(stmts: &[Stmt], entries: &mut Vec) { + for stmt in stmts { + match stmt { + Stmt::FunctionDef(func) => { + maybe_push_test(entries, func.name.as_str(), func.name.range, None, false); + } + Stmt::ClassDef(class_def) => { + let is_unittest = is_unittest_class(class_def); + if is_test_class(class_def, is_unittest) { + entries.push(RunnableEntry { + range: class_def.name.range, + kind: RunnableKind::Test, + test_name: None, + class_name: Some(class_def.name.as_str().to_owned()), + is_unittest, + }); + } + collect_class_entries( + &class_def.body, + entries, + class_def.name.as_str(), + is_unittest, + ); + } + Stmt::If(stmt_if) if Ast::is_main_guard(&stmt_if.test) => { + entries.push(RunnableEntry { + range: stmt_if.range, + kind: RunnableKind::Run, + test_name: None, + class_name: None, + is_unittest: false, + }); + } + _ => {} + } + } +} + +fn collect_class_entries( + stmts: &[Stmt], + entries: &mut Vec, + class_name: &str, + is_unittest: bool, +) { + for stmt in stmts { + if let Stmt::FunctionDef(func) = stmt { + maybe_push_test( + entries, + func.name.as_str(), + func.name.range, + Some(class_name.to_owned()), + is_unittest, + ); + } + } +} + +fn maybe_push_test( + entries: &mut Vec, + name: &str, + range: TextRange, + class_name: Option, + is_unittest: bool, +) { + if is_test_name(name) { + entries.push(RunnableEntry { + range, + kind: RunnableKind::Test, + test_name: Some(name.to_owned()), + class_name, + is_unittest, + }); + } +} + +fn is_test_name(name: &str) -> bool { + name.starts_with("test_") +} + +fn is_test_class(class_def: &StmtClassDef, is_unittest: bool) -> bool { + class_def.name.as_str().starts_with("Test") || is_unittest +} + +fn is_unittest_class(class_def: &StmtClassDef) -> bool { + class_def.bases().iter().any(is_unittest_base) +} + +fn is_unittest_base(base: &Expr) -> bool { + match base { + Expr::Name(name) => name.id.as_str().ends_with("TestCase"), + Expr::Attribute(ExprAttribute { attr, .. }) => attr.id.as_str().ends_with("TestCase"), + _ => false, + } +} diff --git a/pyrefly/lib/test/lsp/lsp_interaction/basic.rs b/pyrefly/lib/test/lsp/lsp_interaction/basic.rs index 0dfdcb2083..3c19c38714 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/basic.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/basic.rs @@ -43,7 +43,7 @@ fn test_initialize_basic() { "codeActionKinds": ["quickfix", "refactor.extract", "refactor.rewrite", "refactor.delete", "refactor.move", "refactor.inline", "source.fixAll", "source.fixAll.pyrefly"] }, "codeLensProvider": { - "resolveProvider": false, + "resolveProvider": true, }, "completionProvider": { "resolveProvider": true, diff --git a/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs b/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs index 2ca23880c9..3ceb39ee19 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs @@ -5,8 +5,14 @@ * LICENSE file in the root directory of this source tree. */ +use std::cell::RefCell; +use std::path::Path; + use lsp_types::CodeLens; +use lsp_types::CodeLensOptions; use lsp_types::CodeLensRequest; +use lsp_types::CodeLensResolveRequest; +use lsp_types::Location; use lsp_types::Uri; use pyrefly_lsp_test::IndexingMode; use pyrefly_lsp_test::LspArgs; @@ -26,6 +32,77 @@ fn runnable_code_lens_config() -> serde_json::Value { }]) } +fn indexed_interaction(root_path: &Path) -> LspInteraction { + let scope_uri = 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_uri)]), + ..Default::default() + }) + .unwrap(); + interaction +} + +fn request_code_lenses(interaction: &LspInteraction, uri: Uri) -> Vec { + let lenses = RefCell::new(Vec::new()); + interaction + .client + .send_request::(json!({ + "textDocument": { + "uri": uri + } + })) + .expect_response_with(|response| { + lenses.replace(response.unwrap_or_default()); + true + }) + .unwrap(); + lenses.into_inner() +} + +fn resolve_code_lens(interaction: &LspInteraction, lens: CodeLens) -> CodeLens { + let resolved = RefCell::new(None); + interaction + .client + .send_request::(serde_json::to_value(lens).unwrap()) + .expect_response_with(|lens| { + resolved.replace(Some(lens)); + true + }) + .unwrap(); + resolved.into_inner().unwrap() +} + +#[test] +fn test_code_lens_advertises_resolve_provider() { + let interaction = LspInteraction::new(); + 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(true), + work_done_progress_options: Default::default(), + }) + }) + .unwrap(); + interaction.client.send_initialized(); + interaction.shutdown().unwrap(); +} + #[test] fn test_code_lens_for_tests_and_main() { let root = get_test_files_root(); @@ -214,69 +291,112 @@ fn test_code_lens_disabled_by_default() { 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 symbols_uri = Url::from_file_path(root_path.join("symbols.py")).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); - interaction - .initialize(InitializeSettings { - workspace_folders: Some(vec![("test".to_owned(), scope_uri)]), - ..Default::default() - }) - .unwrap(); + let symbols_uri = Uri::from_file_path(root_path.join("symbols.py")).unwrap(); + let interaction = indexed_interaction(&root_path); interaction.client.did_open("symbols.py"); interaction.client.did_open("usage.py"); - interaction - .client - .send_request::(json!({ - "textDocument": { - "uri": symbols_uri - } - })) - .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 - }) + let lenses = request_code_lenses(&interaction, symbols_uri); + assert_eq!(lenses.len(), 3); + assert!( + lenses + .iter() + .all(|lens| lens.command.is_none() && lens.data.is_some()) + ); + for (line, title, count) in [ + (0, "3 references", 3), + (1, "2 references", 2), + (4, "2 references", 2), + ] { + let lens = lenses + .iter() + .find(|lens| lens.range.start.line == line) + .unwrap() + .clone(); + let lens = resolve_code_lens(&interaction, lens); + assert!(has_reference_lens(&lens, line, title, count)); + } + + interaction.shutdown().unwrap(); +} + +#[test] +fn test_reference_code_lens_request_from_notebook_is_empty() { + let root = get_test_files_root(); + let root_path = root.path().join("code_lens_references"); + let interaction = indexed_interaction(&root_path); + interaction.open_notebook("notebook.ipynb", vec!["class NotebookSymbol:\n pass"]); + + let lenses = request_code_lenses( + &interaction, + interaction.cell_uri("notebook.ipynb", "cell1"), + ); + assert!(lenses.is_empty()); + + interaction.shutdown().unwrap(); +} + +#[test] +fn test_reference_code_lens_excludes_notebook_references() { + let root = get_test_files_root(); + let root_path = root.path().join("code_lens_references"); + let symbols_uri = Uri::from_file_path(root_path.join("symbols.py")).unwrap(); + let interaction = indexed_interaction(&root_path); + interaction.client.did_open("symbols.py"); + interaction.client.did_open("usage.py"); + interaction.open_notebook( + "notebook.ipynb", + vec!["from symbols import Greeter", "Greeter()"], + ); + let notebook_uris = [ + interaction.cell_uri("notebook.ipynb", "cell1"), + interaction.cell_uri("notebook.ipynb", "cell2"), + ]; + + let lenses = request_code_lenses(&interaction, symbols_uri); + let greeter_lens = lenses + .into_iter() + .find(|lens| lens.range.start.line == 0) .unwrap(); + let resolved = resolve_code_lens(&interaction, greeter_lens); + let locations = reference_locations(&resolved); + assert_eq!(locations.len(), 3); + assert!( + locations + .iter() + .all(|location| !notebook_uris.contains(&location.uri)) + ); interaction.shutdown().unwrap(); } fn has_reference_lens( - lenses: &[CodeLens], + lens: &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; - }; + let Some(command) = &lens.command else { + return false; + }; + lens.range.start.line == line + && command.title == expected_title + && command.command == "editor.action.showReferences" + && reference_locations(lens).len() == expected_locations +} + +fn reference_locations(lens: &CodeLens) -> Vec { + let arguments = lens + .command + .as_ref() + .and_then(|command| command.arguments.as_ref()) + .expect("Resolved reference lens should have command arguments"); + serde_json::from_value( arguments .get(2) - .and_then(|value| value.as_array()) - .is_some_and(|locations| locations.len() == expected_locations) - }) + .expect("Reference locations should be the third argument") + .clone(), + ) + .expect("Reference locations should deserialize") }