Files
ruff/crates/ty_server/tests/e2e/notebook.rs
Andrew Gallant 426125f5c0 [ty] Stabilize auto-import
While still under development, it's far enough along now that we think
it's worth enabling it by default. This should also help give us
feedback for how it behaves.

This PR adds a "completion settings" grouping similar to inlay hints. We
only have an auto-import setting there now, but I expect we'll add more
options to configure completion behavior in the future.

Closes astral-sh/ty#1765
2025-12-09 09:40:38 -05:00

531 lines
16 KiB
Rust

use insta::assert_json_snapshot;
use lsp_types::{NotebookCellKind, Position, Range};
use ruff_db::system::SystemPath;
use ty_server::ClientOptions;
use crate::{TestServer, TestServerBuilder};
static FILTERS: &[(&str, &str)] = &[(r#""sortText": "[0-9 ]+""#, r#""sortText": "[RANKING]""#)];
#[test]
fn publish_diagnostics_open() -> anyhow::Result<()> {
let mut server = TestServerBuilder::new()?
.enable_diagnostic_related_information(true)
.build()
.wait_until_workspaces_are_initialized();
server.initialization_result().unwrap();
let mut builder = NotebookBuilder::virtual_file("test.ipynb");
builder.add_python_cell(
r#"from typing import Literal
type Style = Literal["italic", "bold", "underline"]"#,
);
builder.add_python_cell(
r#"def with_style(line: str, word, style: Style) -> str:
if style == "italic":
return line.replace(word, f"*{word}*")
elif style == "bold":
return line.replace(word, f"__{word}__")
position = line.find(word)
output = line + "\n"
output += " " * position
output += "-" * len(word)
"#,
);
builder.add_python_cell(
r#"print(with_style("ty is a fast type checker for Python.", "fast", "underlined"))
"#,
);
builder.open(&mut server);
let cell1_diagnostics =
server.await_notification::<lsp_types::notification::PublishDiagnostics>();
let cell2_diagnostics =
server.await_notification::<lsp_types::notification::PublishDiagnostics>();
let cell3_diagnostics =
server.await_notification::<lsp_types::notification::PublishDiagnostics>();
assert_json_snapshot!([cell1_diagnostics, cell2_diagnostics, cell3_diagnostics]);
Ok(())
}
#[test]
fn diagnostic_end_of_file() -> anyhow::Result<()> {
let mut server = TestServerBuilder::new()?
.build()
.wait_until_workspaces_are_initialized();
server.initialization_result().unwrap();
let mut builder = NotebookBuilder::virtual_file("test.ipynb");
builder.add_python_cell(
r#"from typing import Literal
type Style = Literal["italic", "bold", "underline"]"#,
);
builder.add_python_cell(
r#"def with_style(line: str, word, style: Style) -> str:
if style == "italic":
return line.replace(word, f"*{word}*")
elif style == "bold":
return line.replace(word, f"__{word}__")
position = line.find(word)
output = line + "\n"
output += " " * position
output += "-" * len(word)
"#,
);
let cell_3 = builder.add_python_cell(
r#"with_style("test", "word", "underline")
IOError"#,
);
let notebook_url = builder.open(&mut server);
server.collect_publish_diagnostic_notifications(3);
server.send_notification::<lsp_types::notification::DidChangeNotebookDocument>(
lsp_types::DidChangeNotebookDocumentParams {
notebook_document: lsp_types::VersionedNotebookDocumentIdentifier {
version: 0,
uri: notebook_url,
},
change: lsp_types::NotebookDocumentChangeEvent {
metadata: None,
cells: Some(lsp_types::NotebookDocumentCellChange {
structure: None,
data: None,
text_content: Some(vec![lsp_types::NotebookDocumentChangeTextContent {
document: lsp_types::VersionedTextDocumentIdentifier {
uri: cell_3,
version: 0,
},
changes: {
vec![lsp_types::TextDocumentContentChangeEvent {
range: Some(Range::new(Position::new(0, 16), Position::new(0, 17))),
range_length: Some(1),
text: String::new(),
}]
},
}]),
}),
},
},
);
let diagnostics = server.collect_publish_diagnostic_notifications(3);
assert_json_snapshot!(diagnostics);
Ok(())
}
#[test]
fn semantic_tokens() -> anyhow::Result<()> {
let mut server = TestServerBuilder::new()?
.build()
.wait_until_workspaces_are_initialized();
server.initialization_result().unwrap();
let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
let first_cell = builder.add_python_cell(
r#"from typing import Literal
type Style = Literal["italic", "bold", "underline"]"#,
);
let second_cell = builder.add_python_cell(
r#"def with_style(line: str, word, style: Style) -> str:
if style == "italic":
return line.replace(word, f"*{word}*")
elif style == "bold":
return line.replace(word, f"__{word}__")
position = line.find(word)
output = line + "\n"
output += " " * position
output += "-" * len(word)
"#,
);
let third_cell = builder.add_python_cell(
r#"print(with_style("ty is a fast type checker for Python.", "fast", "underlined"))
"#,
);
builder.open(&mut server);
let cell1_tokens = semantic_tokens_full_for_cell(&mut server, &first_cell);
let cell2_tokens = semantic_tokens_full_for_cell(&mut server, &second_cell);
let cell3_tokens = semantic_tokens_full_for_cell(&mut server, &third_cell);
assert_json_snapshot!([cell1_tokens, cell2_tokens, cell3_tokens]);
server.collect_publish_diagnostic_notifications(3);
Ok(())
}
#[test]
fn swap_cells() -> anyhow::Result<()> {
let mut server = TestServerBuilder::new()?
.build()
.wait_until_workspaces_are_initialized();
server.initialization_result().unwrap();
let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
let first_cell = builder.add_python_cell(
r#"b = a
"#,
);
let second_cell = builder.add_python_cell(r#"a = 10"#);
builder.add_python_cell(r#"c = b"#);
let notebook = builder.open(&mut server);
let diagnostics = server.collect_publish_diagnostic_notifications(3);
assert_json_snapshot!(diagnostics, @r###"
{
"vscode-notebook-cell://src/test.ipynb#0": [
{
"range": {
"start": {
"line": 0,
"character": 4
},
"end": {
"line": 0,
"character": 5
}
},
"severity": 1,
"code": "unresolved-reference",
"codeDescription": {
"href": "https://ty.dev/rules#unresolved-reference"
},
"source": "ty",
"message": "Name `a` used when not defined"
}
],
"vscode-notebook-cell://src/test.ipynb#1": [],
"vscode-notebook-cell://src/test.ipynb#2": []
}
"###);
// Re-order the cells from `b`, `a`, `c` to `a`, `b`, `c` (swapping cell 1 and 2)
server.send_notification::<lsp_types::notification::DidChangeNotebookDocument>(
lsp_types::DidChangeNotebookDocumentParams {
notebook_document: lsp_types::VersionedNotebookDocumentIdentifier {
version: 1,
uri: notebook,
},
change: lsp_types::NotebookDocumentChangeEvent {
metadata: None,
cells: Some(lsp_types::NotebookDocumentCellChange {
structure: Some(lsp_types::NotebookDocumentCellChangeStructure {
array: lsp_types::NotebookCellArrayChange {
start: 0,
delete_count: 2,
cells: Some(vec![
lsp_types::NotebookCell {
kind: NotebookCellKind::Code,
document: second_cell,
metadata: None,
execution_summary: None,
},
lsp_types::NotebookCell {
kind: NotebookCellKind::Code,
document: first_cell,
metadata: None,
execution_summary: None,
},
]),
},
did_open: None,
did_close: None,
}),
data: None,
text_content: None,
}),
},
},
);
let diagnostics = server.collect_publish_diagnostic_notifications(3);
assert_json_snapshot!(diagnostics, @r###"
{
"vscode-notebook-cell://src/test.ipynb#0": [],
"vscode-notebook-cell://src/test.ipynb#1": [],
"vscode-notebook-cell://src/test.ipynb#2": []
}
"###);
Ok(())
}
#[test]
fn auto_import() -> anyhow::Result<()> {
let mut server = TestServerBuilder::new()?
.with_workspace(
SystemPath::new("src"),
Some(ClientOptions::default().with_auto_import(true)),
)?
.build()
.wait_until_workspaces_are_initialized();
server.initialization_result().unwrap();
let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
builder.add_python_cell(
r#"from typing import TYPE_CHECKING
"#,
);
let second_cell = builder.add_python_cell(
r#"# leading comment
b: Litera
"#,
);
builder.open(&mut server);
server.collect_publish_diagnostic_notifications(2);
let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9));
insta::with_settings!({
filters => FILTERS.iter().copied(),
}, {
assert_json_snapshot!(completions);
});
Ok(())
}
#[test]
fn auto_import_same_cell() -> anyhow::Result<()> {
let mut server = TestServerBuilder::new()?
.with_workspace(
SystemPath::new("src"),
Some(ClientOptions::default().with_auto_import(true)),
)?
.build()
.wait_until_workspaces_are_initialized();
server.initialization_result().unwrap();
let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
let first_cell = builder.add_python_cell(
r#"from typing import TYPE_CHECKING
b: Litera
"#,
);
builder.open(&mut server);
server.collect_publish_diagnostic_notifications(1);
let completions = literal_completions(&mut server, &first_cell, Position::new(1, 9));
insta::with_settings!({
filters => FILTERS.iter().copied(),
}, {
assert_json_snapshot!(completions);
});
Ok(())
}
#[test]
fn auto_import_from_future() -> anyhow::Result<()> {
let mut server = TestServerBuilder::new()?
.with_workspace(
SystemPath::new("src"),
Some(ClientOptions::default().with_auto_import(true)),
)?
.build()
.wait_until_workspaces_are_initialized();
server.initialization_result().unwrap();
let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
builder.add_python_cell(r#"from typing import TYPE_CHECKING"#);
let second_cell = builder.add_python_cell(
r#"from __future__ import annotations
b: Litera
"#,
);
builder.open(&mut server);
server.collect_publish_diagnostic_notifications(2);
let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9));
insta::with_settings!({
filters => FILTERS.iter().copied(),
}, {
assert_json_snapshot!(completions);
});
Ok(())
}
#[test]
fn auto_import_docstring() -> anyhow::Result<()> {
let mut server = TestServerBuilder::new()?
.with_workspace(
SystemPath::new("src"),
Some(ClientOptions::default().with_auto_import(true)),
)?
.build()
.wait_until_workspaces_are_initialized();
server.initialization_result().unwrap();
let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
builder.add_python_cell(
r#"from typing import TYPE_CHECKING
"#,
);
let second_cell = builder.add_python_cell(
r#""""A cell level docstring"""
b: Litera
"#,
);
builder.open(&mut server);
server.collect_publish_diagnostic_notifications(2);
let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9));
insta::with_settings!({
filters => FILTERS.iter().copied(),
}, {
assert_json_snapshot!(completions);
});
Ok(())
}
fn semantic_tokens_full_for_cell(
server: &mut TestServer,
cell_uri: &lsp_types::Url,
) -> Option<lsp_types::SemanticTokensResult> {
let cell1_tokens_req_id = server.send_request::<lsp_types::request::SemanticTokensFullRequest>(
lsp_types::SemanticTokensParams {
work_done_progress_params: lsp_types::WorkDoneProgressParams::default(),
partial_result_params: lsp_types::PartialResultParams::default(),
text_document: lsp_types::TextDocumentIdentifier {
uri: cell_uri.clone(),
},
},
);
server.await_response::<lsp_types::request::SemanticTokensFullRequest>(&cell1_tokens_req_id)
}
#[derive(Debug)]
pub(crate) struct NotebookBuilder {
notebook_url: lsp_types::Url,
// The cells: (cell_metadata, content, language_id)
cells: Vec<(lsp_types::NotebookCell, String, String)>,
}
impl NotebookBuilder {
pub(crate) fn virtual_file(name: &str) -> Self {
let url: lsp_types::Url = format!("vs-code:/{name}").parse().unwrap();
Self {
notebook_url: url,
cells: Vec::new(),
}
}
pub(crate) fn add_python_cell(&mut self, content: &str) -> lsp_types::Url {
let index = self.cells.len();
let id = format!(
"vscode-notebook-cell:/{}#{}",
self.notebook_url.path(),
index
);
let url: lsp_types::Url = id.parse().unwrap();
self.cells.push((
lsp_types::NotebookCell {
kind: NotebookCellKind::Code,
document: url.clone(),
metadata: None,
execution_summary: None,
},
content.to_string(),
"python".to_string(),
));
url
}
pub(crate) fn open(self, server: &mut TestServer) -> lsp_types::Url {
server.send_notification::<lsp_types::notification::DidOpenNotebookDocument>(
lsp_types::DidOpenNotebookDocumentParams {
notebook_document: lsp_types::NotebookDocument {
uri: self.notebook_url.clone(),
notebook_type: "jupyter-notebook".to_string(),
version: 0,
metadata: None,
cells: self.cells.iter().map(|(cell, _, _)| cell.clone()).collect(),
},
cell_text_documents: self
.cells
.iter()
.map(|(cell, content, language_id)| lsp_types::TextDocumentItem {
uri: cell.document.clone(),
language_id: language_id.clone(),
version: 0,
text: content.clone(),
})
.collect(),
},
);
self.notebook_url
}
}
fn literal_completions(
server: &mut TestServer,
cell: &lsp_types::Url,
position: Position,
) -> Vec<lsp_types::CompletionItem> {
let mut items = server.completion_request(cell, position);
// There are a ton of imports we don't care about in here...
// The import bit is that an edit is always restricted to the current cell. That means,
// we can't add `Literal` to the `from typing import TYPE_CHECKING` import in cell 1
items.retain(|item| item.label.starts_with("Litera"));
items
}