Skip to content

Commit b1bf487

Browse files
commtents
1 parent 19e1c14 commit b1bf487

22 files changed

Lines changed: 1169 additions & 280 deletions

File tree

‎crates/pyrefly_lsp_test/src/object_model.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ use pretty_assertions::assert_eq;
7171
use pyrefly::commands::lsp::IndexingMode;
7272
use pyrefly::commands::lsp::LspArgs;
7373
use pyrefly::commands::lsp::run_lsp;
74+
pub use pyrefly::lsp::non_wasm::external_provider::ExternalProvider;
7475
use pyrefly::lsp::non_wasm::external_provider::NoExternalProvider;
7576
use pyrefly::lsp::non_wasm::module_helpers::ThriftRemapper;
7677
use pyrefly::lsp::non_wasm::protocol::JsonRpcMessage;
@@ -1438,6 +1439,7 @@ pub struct LspInteractionArgs {
14381439
pub telemetry: Box<dyn Telemetry>,
14391440
pub thread_count: ThreadCount,
14401441
pub thrift_remapper: Option<ThriftRemapper>,
1442+
pub external_references: Arc<dyn ExternalProvider>,
14411443
}
14421444

14431445
impl Default for LspInteractionArgs {
@@ -1451,6 +1453,7 @@ impl Default for LspInteractionArgs {
14511453
telemetry: Box::new(NoTelemetry),
14521454
thread_count: TEST_THREAD_COUNT,
14531455
thrift_remapper: None,
1456+
external_references: Arc::new(NoExternalProvider),
14541457
}
14551458
}
14561459
}
@@ -1472,6 +1475,7 @@ impl LspInteraction {
14721475
telemetry,
14731476
thread_count,
14741477
thrift_remapper,
1478+
external_references,
14751479
} = args;
14761480

14771481
let ((conn_client, _client_reader), (conn_server, server_reader)) = Connection::memory();
@@ -1489,7 +1493,7 @@ impl LspInteraction {
14891493
None,
14901494
thrift_remapper,
14911495
&*telemetry,
1492-
Arc::new(NoExternalProvider),
1496+
external_references,
14931497
None,
14941498
thread_count,
14951499
);

‎crates/pyrefly_util/src/telemetry.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ pub enum QueueName {
4444
LspQueue,
4545
RecheckQueue,
4646
FindReferenceQueue,
47+
CodeLensQueue,
4748
SourceDbQueue,
4849
}
4950

@@ -59,6 +60,7 @@ impl QueueName {
5960
Self::LspQueue => "lsp_queue",
6061
Self::RecheckQueue => "recheck_queue",
6162
Self::FindReferenceQueue => "find_reference_queue",
63+
Self::CodeLensQueue => "code_lens_queue",
6264
Self::SourceDbQueue => "sourcedb_queue",
6365
}
6466
}

‎lsp/package.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -397,6 +397,12 @@
397397
"description": "Enable Pyrefly's Run/Test CodeLens actions for Python files.",
398398
"scope": "resource"
399399
},
400+
"python.pyrefly.referencesCodeLens": {
401+
"type": "boolean",
402+
"default": false,
403+
"description": "Show reference counts above Python definitions. Requires indexing. Run/Test CodeLens uses a separate setting.",
404+
"scope": "resource"
405+
},
400406
"python.pyrefly.streamDiagnostics": {
401407
"type": "boolean",
402408
"default": true,

‎lsp/src/code-lens.ts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {execFile} from 'child_process';
1111
import * as path from 'path';
1212
import * as vscode from 'vscode';
1313
import {ExtensionContext} from 'vscode';
14+
import type {Location as ProtocolLocation} from 'vscode-languageclient/node';
1415
import {PythonEnvironment} from './python-environment';
1516

1617
type CodeLensPosition = {
@@ -323,6 +324,34 @@ export function registerCodeLensCommands(
323324
context: ExtensionContext,
324325
pythonEnv: PythonEnvironment,
325326
): void {
327+
context.subscriptions.push(
328+
vscode.commands.registerCommand(
329+
'pyrefly.showReferences',
330+
async (
331+
uri: string,
332+
position: CodeLensPosition,
333+
locations: ProtocolLocation[],
334+
) => {
335+
await vscode.commands.executeCommand(
336+
'editor.action.showReferences',
337+
vscode.Uri.parse(uri),
338+
new vscode.Position(position.line, position.character),
339+
locations.map(location =>
340+
new vscode.Location(
341+
vscode.Uri.parse(location.uri),
342+
new vscode.Range(
343+
location.range.start.line,
344+
location.range.start.character,
345+
location.range.end.line,
346+
location.range.end.character,
347+
),
348+
),
349+
),
350+
);
351+
},
352+
),
353+
);
354+
326355
context.subscriptions.push(
327356
vscode.commands.registerCommand('pyrefly.runTest', async args => {
328357
const parsedArgs = parseRunTestArgs(args);

‎lsp/src/extension.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,7 @@ export async function activate(
297297
customHoverProvider: supportsHoverVerbosity,
298298
findFileUsages: true,
299299
pushTypeErrorDisplayStatus: true,
300+
supportsReferencesCodeLens: true,
300301
},
301302
};
302303

‎lsp/src/test/extension.test.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,4 +47,38 @@ suite('Extension Test Suite', () => {
4747
await fs.rm(directory, {recursive: true, force: true});
4848
}
4949
});
50+
51+
test('Show reference CodeLens arguments through the VS Code command', async function () {
52+
this.timeout(30000);
53+
assert.ok(extension);
54+
await extension.activate();
55+
56+
const directory = await fs.mkdtemp(join(tmpdir(), 'pyrefly-references-'));
57+
const uri = vscode.Uri.file(join(directory, 'references.py'));
58+
try {
59+
await vscode.workspace.fs.writeFile(
60+
uri,
61+
Buffer.from('def foo(): pass\nfoo()\n'),
62+
);
63+
const document = await vscode.workspace.openTextDocument(uri);
64+
await vscode.window.showTextDocument(document);
65+
// The built-in command rejects the serialized URI and position without conversion.
66+
await vscode.commands.executeCommand(
67+
'pyrefly.showReferences',
68+
uri.toString(),
69+
{line: 0, character: 4},
70+
[{
71+
uri: uri.toString(),
72+
range: {
73+
start: {line: 1, character: 0},
74+
end: {line: 1, character: 3},
75+
},
76+
}],
77+
);
78+
} finally {
79+
await vscode.commands.executeCommand('closeReferenceSearch');
80+
await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
81+
await fs.rm(directory, {recursive: true, force: true});
82+
}
83+
});
5084
});

‎pyrefly/Cargo.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@ path = "bin/main.rs"
2020
name = "pyrefly_laziness_tests"
2121
path = "test_laziness/mod.rs"
2222

23+
[[bench]]
24+
name = "code_lens"
25+
harness = false
26+
2327
[[bench]]
2428
name = "commit"
2529
harness = false

‎pyrefly/benches/README.md‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,17 @@
1+
# Reference CodeLens benchmark
2+
3+
`code_lens` measures warm LSP requests against a generated 10,000-line Python
4+
file with 2,500 function definitions and ten importing modules. It measures
5+
listing unresolved lenses separately from resolving twenty visible lenses.
6+
Fixture creation, initial analysis, and server shutdown are not timed.
7+
The fixture does not use an external reference provider.
8+
9+
```bash
10+
PYREFLY_LOG=off cargo bench -p pyrefly --bench code_lens -- --noplot
11+
```
12+
13+
These synthetic results do not establish performance for every workspace.
14+
115
# PyTorch benchmarks
216

317
Real-world benchmarks over a large, pinned PyTorch checkout (15k+ Python files)

‎pyrefly/benches/code_lens.rs‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
*
4+
* This source code is licensed under the MIT license found in the
5+
* LICENSE file in the root directory of this source tree.
6+
*/
7+
8+
//! Measure listing and resolving visible reference lenses in a 10,000-line file.
9+
10+
use std::cell::RefCell;
11+
use std::fs;
12+
use std::time::Duration;
13+
14+
use criterion::Criterion;
15+
use criterion::criterion_group;
16+
use criterion::criterion_main;
17+
use lsp_types::CodeLens;
18+
use lsp_types::CodeLensRequest;
19+
use lsp_types::CodeLensResolveRequest;
20+
use lsp_types::Uri;
21+
use pyrefly_lsp_test::IndexingMode;
22+
use pyrefly_lsp_test::LspArgs;
23+
use pyrefly_lsp_test::object_model::InitializeSettings;
24+
use pyrefly_lsp_test::object_model::LspInteraction;
25+
use pyrefly_lsp_test::object_model::LspInteractionArgs;
26+
use pyrefly_util::telemetry::NoTelemetry;
27+
use serde_json::json;
28+
use tempfile::TempDir;
29+
30+
const DEFINITIONS: usize = 2500;
31+
const VISIBLE_LENSES: usize = 20;
32+
const IMPORTERS: usize = 10;
33+
34+
/// Measure warm LSP requests with fixture creation and initial analysis excluded.
35+
fn code_lens(c: &mut Criterion) {
36+
let root = TempDir::new().unwrap();
37+
fs::write(
38+
root.path().join("pyrefly.toml"),
39+
"project-includes = [\"**/*.py\"]\n",
40+
)
41+
.unwrap();
42+
let mut source = String::new();
43+
for index in 0..DEFINITIONS {
44+
source.push_str(&format!(
45+
"def function_{index}() -> int:\n return {index}\n\nfunction_{index}()\n"
46+
));
47+
}
48+
assert_eq!(source.lines().count(), 10_000);
49+
fs::write(root.path().join("large.py"), source).unwrap();
50+
let mut usage = "import large\n".to_owned();
51+
for index in 0..VISIBLE_LENSES {
52+
usage.push_str(&format!("large.function_{index}()\n"));
53+
}
54+
for index in 0..IMPORTERS {
55+
fs::write(root.path().join(format!("usage_{index}.py")), &usage).unwrap();
56+
}
57+
58+
let uri = Uri::from_file_path(root.path().join("large.py")).unwrap();
59+
let mut interaction = LspInteraction::new_with_args(LspInteractionArgs {
60+
args: LspArgs {
61+
indexing_mode: IndexingMode::LazyBlocking,
62+
..LspInteractionArgs::default().args
63+
},
64+
telemetry: Box::new(NoTelemetry),
65+
..Default::default()
66+
});
67+
interaction.client.set_message_logging(false);
68+
interaction.set_root(root.path().to_path_buf());
69+
interaction
70+
.initialize(InitializeSettings {
71+
workspace_folders: Some(vec![(
72+
"benchmark".to_owned(),
73+
Uri::from_file_path(root.path()).unwrap(),
74+
)]),
75+
initialization_options: Some(json!({"pyrefly": {
76+
"supportsReferencesCodeLens": true,
77+
"referencesCodeLens": true
78+
}})),
79+
..Default::default()
80+
})
81+
.unwrap();
82+
interaction.client.did_open("large.py");
83+
let params = json!({"textDocument": {"uri": uri}});
84+
let lenses = RefCell::new(Vec::<CodeLens>::new());
85+
interaction
86+
.client
87+
.send_request::<CodeLensRequest>(params.clone())
88+
.expect_response_with(|response| {
89+
*lenses.borrow_mut() = response.unwrap();
90+
assert_eq!(lenses.borrow().len(), DEFINITIONS);
91+
true
92+
})
93+
.unwrap();
94+
let visible: Vec<_> = lenses
95+
.into_inner()
96+
.into_iter()
97+
.take(VISIBLE_LENSES)
98+
.map(|lens| json!(lens))
99+
.collect();
100+
let resolve_visible = || {
101+
for lens in &visible {
102+
interaction
103+
.client
104+
.send_request::<CodeLensResolveRequest>(lens.clone())
105+
.expect_response_with(|resolved| {
106+
assert_eq!(resolved.command.unwrap().title, "11 references");
107+
true
108+
})
109+
.unwrap();
110+
}
111+
};
112+
resolve_visible();
113+
114+
let mut group = c.benchmark_group("code_lens/10000_lines");
115+
group.sample_size(10);
116+
group.warm_up_time(Duration::from_secs(1));
117+
group.measurement_time(Duration::from_secs(2));
118+
group.bench_function("list", |b| {
119+
b.iter(|| {
120+
interaction
121+
.client
122+
.send_request::<CodeLensRequest>(params.clone())
123+
.expect_response_with(|response| {
124+
let lenses = response.unwrap();
125+
assert_eq!(lenses.len(), DEFINITIONS);
126+
assert!(lenses.iter().all(|lens| lens.command.is_none()));
127+
true
128+
})
129+
.unwrap();
130+
});
131+
});
132+
group.bench_function("resolve_20_visible", |b| b.iter(resolve_visible));
133+
group.finish();
134+
interaction.shutdown().unwrap();
135+
}
136+
137+
criterion_group!(benches, code_lens);
138+
criterion_main!(benches);

‎pyrefly/lib/lsp/non_wasm.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ pub mod move_symbol_new_file;
2020
mod mru;
2121
pub mod protocol;
2222
pub mod queue;
23+
mod runnable;
2324
pub mod safe_delete_file;
2425
pub mod server;
2526
pub mod stdlib;

0 commit comments

Comments
 (0)