fix

Author: asukaminato0721

pyrefly/lib/lsp/non_wasm/server.rs

@@ -40,6 +40,10 @@ use lsp_types::CodeActionParams;
 use lsp_types::CodeActionProviderCapability;
 use lsp_types::CodeActionResponse;
 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;
@@ -152,6 +156,7 @@ use lsp_types::request::CallHierarchyIncomingCalls;
 use lsp_types::request::CallHierarchyOutgoingCalls;
 use lsp_types::request::CallHierarchyPrepare;
 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;
@@ -949,6 +954,12 @@ struct InitializeResult<C> {
     server_info: Option<ServerInfo>,
 }
 
+#[derive(Debug, Clone)]
+struct CodeLensTarget {
+    range: Range,
+    definition: FindDefinitionItemWithDocstring,
+}
+
 pub fn initialize_finish<C: Serialize>(
     sender: &Sender<Message>,
     reader: &mut MessageReader,
@@ -1164,6 +1175,14 @@ pub fn capabilities(
             ]),
             ..Default::default()
         })),
+        code_lens_provider: match indexing_mode {
+            IndexingMode::None => None,
+            IndexingMode::LazyNonBlockingBackground | IndexingMode::LazyBlocking => {
+                Some(CodeLensOptions {
+                    resolve_provider: Some(false),
+                })
+            }
+        },
         completion_provider: Some(CompletionOptions {
             trigger_characters: Some(vec![".".to_owned(), "'".to_owned(), "\"".to_owned()]),
             resolve_provider: Some(true),
@@ -2024,6 +2043,20 @@ impl Server {
                             )),
                         ));
                     }
+                } else if let Some(params) = as_request::<CodeLensRequest>(&x) {
+                    if let Some(params) = self
+                        .extract_request_params_or_send_err_response::<CodeLensRequest>(
+                            params, &x.id,
+                        )
+                    {
+                        self.set_file_stats(params.text_document.uri.clone(), telemetry_event);
+                        self.code_lens(
+                            x.id,
+                            &transaction,
+                            params,
+                            telemetry_event.activity_key.clone(),
+                        );
+                    }
                 } else if let Some(params) = as_request::<WorkspaceSymbolRequest>(&x) {
                     if let Some(params) = self
                         .extract_request_params_or_send_err_response::<WorkspaceSymbolRequest>(
@@ -4259,6 +4292,61 @@ impl Server {
         (!actions.is_empty()).then_some(actions)
     }
 
+    fn code_lens_targets(
+        &self,
+        transaction: &Transaction<'_>,
+        handle: &Handle,
+        uri: &Url,
+    ) -> Option<Vec<CodeLensTarget>> {
+        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)?;
+        let symbols = transaction.symbols(handle)?;
+        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 Some(definition) = transaction
+                .find_definition(
+                    handle,
+                    position,
+                    FindPreference {
+                        import_behavior: ImportBehavior::StopAtRenamedImports,
+                        ..Default::default()
+                    },
+                )
+                .into_iter()
+                .next()
+            else {
+                continue;
+            };
+            let key = (definition.module.path().dupe(), definition.definition_range);
+            if !seen.insert(key) {
+                continue;
+            }
+            targets.push(CodeLensTarget {
+                range: symbol.selection_range,
+                definition,
+            });
+        }
+        Some(targets)
+    }
+
     fn document_highlight(
         &self,
         transaction: &Transaction<'_>,
@@ -4508,6 +4596,113 @@ impl Server {
         );
     }
 
+    fn code_lens<'a>(
+        &'a self,
+        request_id: RequestId,
+        transaction: &Transaction<'a>,
+        params: CodeLensParams,
+        activity_key: Option<ActivityKey>,
+    ) {
+        let uri = &params.text_document.uri;
+        if self.open_notebook_cells.read().contains_key(uri) {
+            return self.send_response(new_response(request_id, Ok(Some(Vec::<CodeLens>::new()))));
+        }
+        let Some(handle) = self.make_handle_if_enabled(uri, Some(CodeLensRequest::METHOD)) else {
+            return self.send_response(new_response::<Option<Vec<CodeLens>>>(request_id, Ok(None)));
+        };
+        let Some(targets) = self.code_lens_targets(transaction, &handle, uri) else {
+            return self.send_response(new_response::<Option<Vec<CodeLens>>>(request_id, Ok(None)));
+        };
+        if targets.is_empty() {
+            return self.send_response(new_response(request_id, Ok(Some(Vec::<CodeLens>::new()))));
+        }
+
+        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 = Vec::new();
+                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)),
+                )));
+            }),
+        );
+    }
+
     fn rename<'a>(
         &'a self,
         request_id: RequestId,

pyrefly/lib/lsp/non_wasm/workspace.rs

@@ -215,6 +215,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,
@@ -242,6 +244,7 @@ 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,

pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs

@@ -0,0 +1,97 @@
+/*
+ * 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::CodeLensOptions;
+use lsp_types::Url;
+use pyrefly::commands::lsp::IndexingMode;
+
+use crate::object_model::InitializeSettings;
+use crate::object_model::LspInteraction;
+use crate::util::get_test_files_root;
+
+#[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)
+    })
+}

pyrefly/lib/test/lsp/lsp_interaction/mod.rs

@@ -13,6 +13,7 @@ mod init;
 
 mod basic;
 mod call_hierarchy;
+mod code_lens;
 mod completion;
 mod configuration;
 mod convert_module_package;

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;
@@ -518,6 +519,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,

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()