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/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 47a5d3df8b..23e08a2da2 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -5,6 +5,7 @@ * LICENSE file in the root directory of this source tree. */ +mod code_lens; mod event_loop; use std::cmp::min; @@ -52,10 +53,9 @@ 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::CodeLensResolveRequest; use lsp_types::CompletionItem; use lsp_types::CompletionList; use lsp_types::CompletionOptions; @@ -299,7 +299,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; @@ -1672,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 { @@ -2376,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); @@ -2700,18 +2700,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::( @@ -2757,6 +2745,37 @@ 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); + 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, + ¶ms, + telemetry_event.activity_key.clone(), + ) + { + 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 .extract_request_params_or_send_err_response::( @@ -3190,24 +3209,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)) } @@ -6019,36 +6020,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.as_str())) - .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/server/code_lens.rs b/pyrefly/lib/lsp/non_wasm/server/code_lens.rs new file mode 100644 index 0000000000..4c383cf376 --- /dev/null +++ b/pyrefly/lib/lsp/non_wasm/server/code_lens.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. + */ + +//! 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::CodeLensRequest; +use lsp_types::Request as _; +use pyrefly_python::module::TextRangeWithModule; +use pyrefly_util::telemetry::ActivityKey; +use pyrefly_util::telemetry::EmptyResponseReason; + +use super::Server; +use crate::commands::lsp::IndexingMode; +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::lsp::ReferenceOptions; +use crate::state::state::Transaction; + +impl Server { + /// Return eager runnable lenses and unresolved reference-count lenses for a Python file. + pub(super) fn code_lens( + &self, + transaction: &Transaction<'_>, + params: CodeLensParams, + ) -> Result>, EmptyResponseReason> { + let uri = ¶ms.text_document.uri; + if self.open_notebook_cells.read().contains_key(uri) { + return Ok(Some(Vec::new())); + } + + 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)) + } + + /// 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, + ErrorCode::InvalidParams as i32, + message, + ))); + return Ok(()); + } + }; + if self.open_notebook_cells.read().contains_key(&data.uri) { + return Err(EmptyResponseReason::NotebookNotSupported); + } + + 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/code_lens.rs b/pyrefly/lib/lsp/non_wasm/server/code_lens/runnable.rs similarity index 59% rename from pyrefly/lib/lsp/non_wasm/code_lens.rs rename to pyrefly/lib/lsp/non_wasm/server/code_lens/runnable.rs index 594358e72b..6084caa3ed 100644 --- a/pyrefly/lib/lsp/non_wasm/code_lens.rs +++ b/pyrefly/lib/lsp/non_wasm/server/code_lens/runnable.rs @@ -5,12 +5,18 @@ * 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; @@ -18,31 +24,78 @@ 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)] -pub enum CodeLensKind { +enum RunnableKind { 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, +struct RunnableEntry { + range: TextRange, + kind: RunnableKind, + test_name: Option, + class_name: Option, + is_unittest: bool, } -pub fn runnable_lsp_code_lens( +/// 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: CodeLensEntry, + entry: RunnableEntry, cwd: Option<&str>, ) -> CodeLens { let (title, command, arguments) = match entry.kind { - CodeLensKind::Run => { + 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 { @@ -50,7 +103,7 @@ pub fn runnable_lsp_code_lens( } ("Run", "pyrefly.runMain", Some(vec![Value::Object(args)])) } - CodeLensKind::Test => { + 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 { @@ -89,24 +142,23 @@ pub fn runnable_lsp_code_lens( } } -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) +/// 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) { +fn collect_module_entries(stmts: &[Stmt], entries: &mut Vec) { for stmt in stmts { match stmt { Stmt::FunctionDef(func) => { @@ -115,9 +167,9 @@ fn collect_module_entries(stmts: &[Stmt], entries: &mut Vec) { Stmt::ClassDef(class_def) => { let is_unittest = is_unittest_class(class_def); if is_test_class(class_def, is_unittest) { - entries.push(CodeLensEntry { + entries.push(RunnableEntry { range: class_def.name.range, - kind: CodeLensKind::Test, + kind: RunnableKind::Test, test_name: None, class_name: Some(class_def.name.as_str().to_owned()), is_unittest, @@ -131,9 +183,9 @@ fn collect_module_entries(stmts: &[Stmt], entries: &mut Vec) { ); } Stmt::If(stmt_if) if Ast::is_main_guard(&stmt_if.test) => { - entries.push(CodeLensEntry { + entries.push(RunnableEntry { range: stmt_if.range, - kind: CodeLensKind::Run, + kind: RunnableKind::Run, test_name: None, class_name: None, is_unittest: false, @@ -146,7 +198,7 @@ fn collect_module_entries(stmts: &[Stmt], entries: &mut Vec) { fn collect_class_entries( stmts: &[Stmt], - entries: &mut Vec, + entries: &mut Vec, class_name: &str, is_unittest: bool, ) { @@ -164,16 +216,16 @@ fn collect_class_entries( } fn maybe_push_test( - entries: &mut Vec, + entries: &mut Vec, name: &str, range: TextRange, class_name: Option, is_unittest: bool, ) { if is_test_name(name) { - entries.push(CodeLensEntry { + entries.push(RunnableEntry { range, - kind: CodeLensKind::Test, + kind: RunnableKind::Test, test_name: Some(name.to_owned()), class_name, is_unittest, @@ -186,10 +238,7 @@ fn is_test_name(name: &str) -> bool { } fn is_test_class(class_def: &StmtClassDef, is_unittest: bool) -> bool { - if class_def.name.as_str().starts_with("Test") { - return true; - } - is_unittest + class_def.name.as_str().starts_with("Test") || is_unittest } fn is_unittest_class(class_def: &StmtClassDef) -> bool { 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 f177a31669..3ceb39ee19 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs @@ -5,11 +5,20 @@ * 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; use pyrefly_lsp_test::object_model::InitializeSettings; use pyrefly_lsp_test::object_model::LspInteraction; +use pyrefly_lsp_test::object_model::LspInteractionArgs; use serde_json::Value; use serde_json::json; @@ -23,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(); @@ -206,3 +286,117 @@ fn test_code_lens_disabled_by_default() { 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 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"); + + 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( + lens: &CodeLens, + line: u32, + expected_title: &str, + expected_locations: usize, +) -> bool { + 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) + .expect("Reference locations should be the third argument") + .clone(), + ) + .expect("Reference locations should deserialize") +} 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)