diff --git a/pyrefly/lib/lsp/non_wasm/document_symbols.rs b/pyrefly/lib/lsp/non_wasm/document_symbols.rs index 805743c087..e60c447c0b 100644 --- a/pyrefly/lib/lsp/non_wasm/document_symbols.rs +++ b/pyrefly/lib/lsp/non_wasm/document_symbols.rs @@ -287,3 +287,42 @@ fn recurse_stmt_adding_symbols<'a>( }; symbols.append(&mut recursed_symbols); } + +pub fn flatten_to_symbol_information( + symbols: Vec, + uri: &lsp_types::Url, +) -> Vec { + let mut results = Vec::new(); + flatten_recursive(symbols, uri, None, &mut results); + results +} + +fn flatten_recursive( + symbols: Vec, + uri: &lsp_types::Url, + container_name: Option, + result: &mut Vec, +) { + for sym in symbols { + let children = sym.children.unwrap_or_default(); + let qualified_name = match &container_name { + Some(parent) => format!("{}.{}", parent, sym.name), + None => sym.name.clone(), + }; + + #[allow(deprecated)] + result.push(lsp_types::SymbolInformation { + name: sym.name, + kind: sym.kind, + tags: sym.tags, + deprecated: sym.deprecated, + location: lsp_types::Location { + uri: uri.clone(), + range: sym.range, + }, + container_name: container_name.clone(), + }); + + flatten_recursive(children, uri, Some(qualified_name), result); + } +} diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 9cdf981f61..31249c1503 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -65,7 +65,6 @@ use lsp_types::DocumentDiagnosticReport; use lsp_types::DocumentHighlight; use lsp_types::DocumentHighlightKind; use lsp_types::DocumentHighlightParams; -use lsp_types::DocumentSymbol; use lsp_types::DocumentSymbolParams; use lsp_types::DocumentSymbolResponse; use lsp_types::FileEvent; @@ -270,6 +269,7 @@ 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; use crate::lsp::non_wasm::external_provider::compute_qualified_name; use crate::lsp::non_wasm::lsp::apply_change_events; @@ -2229,14 +2229,13 @@ impl Server { params, &x.id, ) { - let response = - match self.hierarchical_document_symbols(&transaction, params) { - Ok(response) => response.map(DocumentSymbolResponse::Nested), - Err(reason) => { - telemetry_event.set_empty_response_reason(reason); - None - } - }; + let response = match self.document_symbol(&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) { @@ -5101,34 +5100,38 @@ impl Server { }))) } - fn hierarchical_document_symbols( + fn document_symbol( &self, transaction: &Transaction<'_>, params: DocumentSymbolParams, - ) -> Result>, EmptyResponseReason> { + ) -> Result, EmptyResponseReason> { let uri = ¶ms.text_document.uri; let maybe_cell_idx = self.maybe_get_code_cell_index(uri); let path = self .path_for_uri_or_notebook_cell(uri) .ok_or(EmptyResponseReason::NoFilePath)?; - if self - .workspaces - .get_with(path, |(_, workspace)| workspace.disable_language_services) - { + if self.workspaces.get_with(path, |(_, workspace)| { + workspace.disabled_language_services.is_some() + }) { return Err(EmptyResponseReason::LanguageServicesDisabled); } - let Some(true) = self + let handle = self.make_handle_if_enabled(uri, Some(DocumentSymbolRequest::METHOD))?; + let symbols = transaction.symbols(&handle, maybe_cell_idx); + let supports_hierarchical = self .initialize_params .capabilities .text_document .as_ref() .and_then(|t| t.document_symbol.as_ref()) .and_then(|d| d.hierarchical_document_symbol_support) - else { - return Ok(None); - }; - let handle = self.make_handle_if_enabled(uri, Some(DocumentSymbolRequest::METHOD))?; - Ok(transaction.symbols(&handle, maybe_cell_idx)) + == Some(true); + Ok(symbols.map(|syms| { + if supports_hierarchical { + DocumentSymbolResponse::Nested(syms) + } else { + DocumentSymbolResponse::Flat(flatten_to_symbol_information(syms, uri)) + } + })) } /// Run local and external workspace symbol queries in parallel, merging diff --git a/pyrefly/lib/test/lsp/document_symbols.rs b/pyrefly/lib/test/lsp/document_symbols.rs index 8b349b35db..59b9f47321 100644 --- a/pyrefly/lib/test/lsp/document_symbols.rs +++ b/pyrefly/lib/test/lsp/document_symbols.rs @@ -8,6 +8,7 @@ use pretty_assertions::assert_eq; use pyrefly_build::handle::Handle; +use crate::lsp::non_wasm::document_symbols::flatten_to_symbol_information; use crate::state::state::State; use crate::test::util::get_batched_lsp_operations_report_no_cursor; use crate::test::util::get_batched_lsp_operations_report_no_cursor_allow_error; @@ -21,6 +22,17 @@ fn get_test_report(state: &State, handle: &Handle) -> String { } } +fn get_flat_symbol_report(state: &State, handle: &Handle) -> String { + let transactions = state.transaction(); + let uri = lsp_types::Url::parse("file:///main.py").unwrap(); + if let Some(symbols) = transactions.symbols(handle, None) { + let flat = flatten_to_symbol_information(symbols, &uri); + serde_json::to_string_pretty(&flat).unwrap() + } else { + "No document symbols found".to_owned() + } +} + #[test] fn function_test() { let code = r#" @@ -1169,3 +1181,89 @@ b = 2 assert_eq!(config_children[0].name, "DEBUG"); assert_eq!(config_children[1].name, "b"); } + +#[test] +fn test_flatten_basic_class() { + let code = r#" +class MyClass: + def __init__(self): + self.x = 1 + + def method1(self): + return self.x + + def method2(self, y): + return self.x + y +"#; + + let report = + get_batched_lsp_operations_report_no_cursor(&[("main", code)], get_flat_symbol_report); + let uri = lsp_types::Url::parse("file:///main.py").unwrap(); + + #[allow(deprecated)] + let flat: Vec = + serde_json::from_str(&report.split("\n").skip(2).collect::>().join("\n")).unwrap(); + + assert_eq!(flat.len(), 4); + + assert_eq!(flat[0].name, "MyClass"); + assert_eq!(flat[0].container_name, None); + assert_eq!(flat[0].location.uri, uri); + + assert_eq!(flat[1].name, "__init__"); + assert_eq!(flat[1].container_name, Some("MyClass".to_owned())); + + assert_eq!(flat[2].name, "method1"); + assert_eq!(flat[2].container_name, Some("MyClass".to_owned())); + + assert_eq!(flat[3].name, "method2"); + assert_eq!(flat[3].container_name, Some("MyClass".to_owned())); +} + +#[test] +fn test_flatten_deep_nesting() { + let code = r#" +class Outer: + class Inner: + def foo(self): + pass +"#; + + let report = + get_batched_lsp_operations_report_no_cursor(&[("main", code)], get_flat_symbol_report); + let uri = lsp_types::Url::parse("file:///main.py").unwrap(); + + #[allow(deprecated)] + let flat: Vec = + serde_json::from_str(&report.split("\n").skip(2).collect::>().join("\n")).unwrap(); + + assert_eq!(flat.len(), 3); + + assert_eq!(flat[0].name, "Outer"); + assert_eq!(flat[0].container_name, None); + assert_eq!(flat[0].location.uri, uri); + + assert_eq!(flat[1].name, "Inner"); + assert_eq!(flat[1].container_name, Some("Outer".to_owned())); + + assert_eq!(flat[2].name, "foo"); + assert_eq!(flat[2].container_name, Some("Outer.Inner".to_owned())); +} + +#[test] +fn test_flatten_empty_file() { + let code = r#" +# just a comment, no symbols +"#; + + let report = + get_batched_lsp_operations_report_no_cursor(&[("main", code)], get_flat_symbol_report); + + // The report is either "No document symbols found" or an empty JSON array + let body = report.split("\n").skip(2).collect::>().join("\n"); + if body.trim() != "No document symbols found" { + #[allow(deprecated)] + let flat: Vec = serde_json::from_str(body.trim()).unwrap(); + assert!(flat.is_empty()); + } +}