Fix document symbol search for editors without hierarchical support

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<DocumentSymbol>,
+    uri: &lsp_types::Url,
+) -> Vec<lsp_types::SymbolInformation> {
+    let mut results = Vec::new();
+    flatten_recursive(symbols, uri, None, &mut results);
+    results
+}
+
+fn flatten_recursive(
+    symbols: Vec<DocumentSymbol>,
+    uri: &lsp_types::Url,
+    container_name: Option<String>,
+    result: &mut Vec<lsp_types::SymbolInformation>,
+) {
+    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);
+    }
+}

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::<WorkspaceSymbolRequest>(&x) {
@@ -5101,34 +5100,38 @@ impl Server {
         })))
     }
 
-    fn hierarchical_document_symbols(
+    fn document_symbol(
         &self,
         transaction: &Transaction<'_>,
         params: DocumentSymbolParams,
-    ) -> Result<Option<Vec<DocumentSymbol>>, EmptyResponseReason> {
+    ) -> Result<Option<DocumentSymbolResponse>, EmptyResponseReason> {
         let uri = &params.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

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<lsp_types::SymbolInformation> =
+        serde_json::from_str(&report.split("\n").skip(2).collect::<Vec<_>>().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<lsp_types::SymbolInformation> =
+        serde_json::from_str(&report.split("\n").skip(2).collect::<Vec<_>>().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::<Vec<_>>().join("\n");
+    if body.trim() != "No document symbols found" {
+        #[allow(deprecated)]
+        let flat: Vec<lsp_types::SymbolInformation> = serde_json::from_str(body.trim()).unwrap();
+        assert!(flat.is_empty());
+    }
+}