Skip to content

Commit 66ad6a1

Browse files
asukaminato0721asukaminato0721
committed
fix
1 parent 947a3ee commit 66ad6a1

6 files changed

Lines changed: 334 additions & 45 deletions

File tree

‎pyrefly/lib/lsp/non_wasm/server.rs‎

Lines changed: 226 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ use lsp_types::CodeActionTriggerKind;
4747
use lsp_types::CodeLens;
4848
use lsp_types::CodeLensOptions;
4949
use lsp_types::CodeLensParams;
50+
use lsp_types::Command;
5051
use lsp_types::CompletionItem;
5152
use lsp_types::CompletionList;
5253
use lsp_types::CompletionOptions;
@@ -1101,6 +1102,12 @@ struct InitializeResult<C> {
11011102
server_info: Option<ServerInfo>,
11021103
}
11031104

1105+
#[derive(Debug, Clone)]
1106+
struct CodeLensTarget {
1107+
range: Range,
1108+
definition: FindDefinitionItemWithDocstring,
1109+
}
1110+
11041111
pub fn initialize_finish<C: Serialize>(
11051112
sender: &Sender<Message>,
11061113
reader: &mut MessageReader,
@@ -2251,18 +2258,6 @@ impl Server {
22512258
};
22522259
self.send_response(new_response(x.id, Ok(response)));
22532260
}
2254-
} else if let Some(params) = as_request::<CodeLensRequest>(&x) {
2255-
if let Some(params) = self
2256-
.extract_request_params_or_send_err_response::<CodeLensRequest>(
2257-
params, &x.id,
2258-
)
2259-
{
2260-
self.set_file_stats(params.text_document.uri.clone(), telemetry_event);
2261-
self.send_response(new_response(
2262-
x.id,
2263-
Ok(self.code_lens(&transaction, params).unwrap_or_default()),
2264-
));
2265-
}
22662261
} else if let Some(params) = as_request::<SemanticTokensFullRequest>(&x) {
22672262
if let Some(params) = self
22682263
.extract_request_params_or_send_err_response::<SemanticTokensFullRequest>(
@@ -2309,6 +2304,23 @@ impl Server {
23092304
};
23102305
self.send_response(new_response(x.id, Ok(response)));
23112306
}
2307+
} else if let Some(params) = as_request::<CodeLensRequest>(&x) {
2308+
if let Some(params) = self
2309+
.extract_request_params_or_send_err_response::<CodeLensRequest>(
2310+
params, &x.id,
2311+
)
2312+
{
2313+
self.set_file_stats(params.text_document.uri.clone(), telemetry_event);
2314+
if let Err(reason) = self.code_lens(
2315+
x.id.clone(),
2316+
&transaction,
2317+
params,
2318+
telemetry_event.activity_key.clone(),
2319+
) {
2320+
self.send_response(new_response(x.id, Ok(None::<()>)));
2321+
telemetry_event.set_empty_response_reason(reason);
2322+
}
2323+
}
23122324
} else if let Some(params) = as_request::<WorkspaceSymbolRequest>(&x) {
23132325
if let Some(params) = self
23142326
.extract_request_params_or_send_err_response::<WorkspaceSymbolRequest>(
@@ -4633,6 +4645,64 @@ impl Server {
46334645
Ok((!actions.is_empty()).then_some(actions))
46344646
}
46354647

4648+
fn code_lens_targets(
4649+
&self,
4650+
transaction: &Transaction<'_>,
4651+
handle: &Handle,
4652+
uri: &Url,
4653+
) -> Result<Vec<CodeLensTarget>, EmptyResponseReason> {
4654+
fn recurse_symbols<'a>(symbols: &'a [DocumentSymbol], out: &mut Vec<&'a DocumentSymbol>) {
4655+
for symbol in symbols {
4656+
if matches!(
4657+
symbol.kind,
4658+
SymbolKind::CLASS | SymbolKind::FUNCTION | SymbolKind::METHOD
4659+
) {
4660+
out.push(symbol);
4661+
}
4662+
if let Some(children) = symbol.children.as_deref() {
4663+
recurse_symbols(children, out);
4664+
}
4665+
}
4666+
}
4667+
4668+
let module_info = transaction
4669+
.get_module_info(handle)
4670+
.ok_or(EmptyResponseReason::ModuleInfoNotFound)?;
4671+
let symbols = transaction
4672+
.symbols(handle, None)
4673+
.ok_or(EmptyResponseReason::ModuleInfoNotFound)?;
4674+
let mut symbol_defs = Vec::new();
4675+
recurse_symbols(&symbols, &mut symbol_defs);
4676+
4677+
let mut seen = SmallSet::new();
4678+
let mut targets = Vec::new();
4679+
for symbol in symbol_defs {
4680+
let position = self.from_lsp_position(uri, &module_info, symbol.selection_range.start);
4681+
let definition = match transaction.find_definition(
4682+
handle,
4683+
position,
4684+
FindPreference {
4685+
import_behavior: ImportBehavior::StopAtRenamedImports,
4686+
..Default::default()
4687+
},
4688+
) {
4689+
Ok(definitions) => definitions.into_vec().swap_remove(0),
4690+
Err(_) => {
4691+
continue;
4692+
}
4693+
};
4694+
let key = (definition.module.path().dupe(), definition.definition_range);
4695+
if !seen.insert(key) {
4696+
continue;
4697+
}
4698+
targets.push(CodeLensTarget {
4699+
range: symbol.selection_range,
4700+
definition,
4701+
});
4702+
}
4703+
Ok(targets)
4704+
}
4705+
46364706
fn document_highlight(
46374707
&self,
46384708
transaction: &Transaction<'_>,
@@ -4884,6 +4954,150 @@ impl Server {
48844954
)
48854955
}
48864956

4957+
fn code_lens<'a>(
4958+
&'a self,
4959+
request_id: RequestId,
4960+
transaction: &Transaction<'a>,
4961+
params: CodeLensParams,
4962+
activity_key: Option<ActivityKey>,
4963+
) -> Result<(), EmptyResponseReason> {
4964+
let uri = &params.text_document.uri;
4965+
if self.open_notebook_cells.read().contains_key(uri) {
4966+
self.send_response(new_response(request_id, Ok(Some(Vec::<CodeLens>::new()))));
4967+
return Ok(());
4968+
}
4969+
let handle = self.make_handle_if_enabled(uri, Some(CodeLensRequest::METHOD))?;
4970+
let runnable_lenses = self.runnable_code_lenses(transaction, &handle, &params)?;
4971+
let targets = if self.indexing_mode == IndexingMode::None {
4972+
Vec::new()
4973+
} else {
4974+
self.code_lens_targets(transaction, &handle, uri)?
4975+
};
4976+
if targets.is_empty() {
4977+
self.send_response(new_response(request_id, Ok(Some(runnable_lenses))));
4978+
return Ok(());
4979+
}
4980+
4981+
let path_remapper = self.path_remapper.clone();
4982+
let source_uri = uri.clone();
4983+
self.find_reference_queue.queue_task(
4984+
TelemetryEventKind::FindFromDefinition,
4985+
Box::new(move |server, _telemetry, telemetry_event| {
4986+
telemetry_event.set_activity_key(activity_key);
4987+
let mut transaction = server.state.cancellable_transaction();
4988+
server
4989+
.cancellation_handles
4990+
.lock()
4991+
.insert(request_id.clone(), transaction.get_cancellation_handle());
4992+
server.validate_in_memory_for_transaction(
4993+
transaction.as_mut(),
4994+
telemetry_event,
4995+
None,
4996+
);
4997+
4998+
let mut lenses = runnable_lenses;
4999+
for target in targets {
5000+
let local_results = match transaction.find_global_references_from_definition(
5001+
*handle.sys_info(),
5002+
target.definition.metadata,
5003+
TextRangeWithModule::new(
5004+
target.definition.module.clone(),
5005+
target.definition.definition_range,
5006+
),
5007+
false,
5008+
) {
5009+
Ok(results) => results,
5010+
Err(Cancelled) => {
5011+
let message = format!("Request {request_id} is canceled");
5012+
info!("{message}");
5013+
server.connection.send(Message::Response(Response::new_err(
5014+
request_id,
5015+
ErrorCode::RequestCanceled as i32,
5016+
message,
5017+
)));
5018+
return;
5019+
}
5020+
};
5021+
5022+
let mut locations = Vec::new();
5023+
for (info, ranges) in local_results {
5024+
if let Some(uri) = module_info_to_uri(&info, path_remapper.as_ref()) {
5025+
for range in ranges {
5026+
locations.push(Location {
5027+
uri: uri.clone(),
5028+
range: info.to_lsp_range(range),
5029+
});
5030+
}
5031+
}
5032+
}
5033+
5034+
let reference_count = locations.len();
5035+
let title = if reference_count == 1 {
5036+
"1 reference".to_owned()
5037+
} else {
5038+
format!("{reference_count} references")
5039+
};
5040+
lenses.push(CodeLens {
5041+
range: target.range,
5042+
command: Some(Command {
5043+
title,
5044+
command: "editor.action.showReferences".to_owned(),
5045+
arguments: Some(vec![
5046+
serde_json::to_value(&source_uri)
5047+
.expect("URI should serialize for code lens"),
5048+
serde_json::to_value(target.range.start)
5049+
.expect("Position should serialize for code lens"),
5050+
serde_json::to_value(&locations)
5051+
.expect("Locations should serialize for code lens"),
5052+
]),
5053+
}),
5054+
data: None,
5055+
});
5056+
}
5057+
5058+
server.cancellation_handles.lock().remove(&request_id);
5059+
server.connection.send(Message::Response(new_response(
5060+
request_id,
5061+
Ok(Some(lenses)),
5062+
)));
5063+
}),
5064+
);
5065+
Ok(())
5066+
}
5067+
5068+
fn runnable_code_lenses(
5069+
&self,
5070+
transaction: &Transaction<'_>,
5071+
handle: &Handle,
5072+
params: &CodeLensParams,
5073+
) -> Result<Vec<CodeLens>, EmptyResponseReason> {
5074+
let uri = &params.text_document.uri;
5075+
let path = self
5076+
.path_for_uri(uri)
5077+
.ok_or(EmptyResponseReason::NoFilePath)?;
5078+
let runnable_code_lens = self
5079+
.workspaces
5080+
.get_with(path.clone(), |(_, workspace)| workspace.runnable_code_lens);
5081+
let maybe_cell_idx = self.maybe_get_cell_index(uri);
5082+
let info = transaction
5083+
.get_module_info(handle)
5084+
.ok_or(EmptyResponseReason::ModuleInfoNotFound)?;
5085+
let entries = transaction
5086+
.runnable_code_lens_entries(handle, uri, runnable_code_lens)
5087+
.ok_or(EmptyResponseReason::ModuleInfoNotFound)?;
5088+
let cwd = self.runnable_code_lens_cwd(&path);
5089+
5090+
let mut lenses = Vec::new();
5091+
for entry in entries {
5092+
if info.to_cell_for_lsp(entry.range.start()) != maybe_cell_idx {
5093+
continue;
5094+
}
5095+
let range = info.to_lsp_range(entry.range);
5096+
lenses.push(runnable_lsp_code_lens(uri, range, entry, cwd.as_deref()));
5097+
}
5098+
Ok(lenses)
5099+
}
5100+
48875101
fn rename<'a>(
48885102
&'a self,
48895103
request_id: RequestId,
@@ -5048,36 +5262,6 @@ impl Server {
50485262
Ok(Some(res))
50495263
}
50505264

5051-
fn code_lens(
5052-
&self,
5053-
transaction: &Transaction<'_>,
5054-
params: CodeLensParams,
5055-
) -> Option<Vec<CodeLens>> {
5056-
let uri = &params.text_document.uri;
5057-
let path = self.path_for_uri(uri)?;
5058-
let runnable_code_lens = self
5059-
.workspaces
5060-
.get_with(path.clone(), |(_, workspace)| workspace.runnable_code_lens);
5061-
let maybe_cell_idx = self.maybe_get_cell_index(uri);
5062-
let handle = self
5063-
.make_handle_if_enabled(uri, Some(CodeLensRequest::METHOD))
5064-
.ok()?;
5065-
let info = transaction.get_module_info(&handle)?;
5066-
let entries = transaction.runnable_code_lens_entries(&handle, uri, runnable_code_lens)?;
5067-
let cwd = self.runnable_code_lens_cwd(&path);
5068-
5069-
let mut lenses = Vec::new();
5070-
for entry in entries {
5071-
if info.to_cell_for_lsp(entry.range.start()) != maybe_cell_idx {
5072-
continue;
5073-
}
5074-
let range = info.to_lsp_range(entry.range);
5075-
lenses.push(runnable_lsp_code_lens(uri, range, entry, cwd.as_deref()));
5076-
}
5077-
5078-
Some(lenses)
5079-
}
5080-
50815265
fn semantic_tokens_full(
50825266
&self,
50835267
transaction: &Transaction<'_>,

‎pyrefly/lib/lsp/non_wasm/workspace.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,8 @@ pub struct DisabledLanguageServices {
238238
#[serde(default)]
239239
pub references: bool,
240240
#[serde(default)]
241+
pub code_lens: bool,
242+
#[serde(default)]
241243
pub rename: bool,
242244
#[serde(default)]
243245
pub signature_help: bool,
@@ -248,8 +250,6 @@ pub struct DisabledLanguageServices {
248250
#[serde(default)]
249251
pub document_symbol: bool,
250252
#[serde(default)]
251-
pub code_lens: bool,
252-
#[serde(default)]
253253
pub semantic_tokens: bool,
254254
#[serde(default)]
255255
pub implementation: bool,
@@ -267,12 +267,12 @@ impl DisabledLanguageServices {
267267
"textDocument/completion" => self.completion,
268268
"textDocument/documentHighlight" => self.document_highlight,
269269
"textDocument/references" => self.references,
270+
"textDocument/codeLens" => self.code_lens,
270271
"textDocument/rename" => self.rename,
271272
"textDocument/signatureHelp" => self.signature_help,
272273
"textDocument/hover" => self.hover,
273274
"textDocument/inlayHint" => self.inlay_hint,
274275
"textDocument/documentSymbol" => self.document_symbol,
275-
"textDocument/codeLens" => self.code_lens,
276276
"textDocument/semanticTokens/full" | "textDocument/semanticTokens/range" => {
277277
self.semantic_tokens
278278
}

0 commit comments

Comments
 (0)