mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-11-20 04:29:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

[ty] Support for notebooks in VS Code (#21175)

Browse source
This commit is contained in:
Micha Reiser 2025-11-13 13:23:19 +01:00 • committed by GitHub
parent d64b2f747c
commit 12e74ae894
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
47 changed files with 1985 additions and 784 deletions
Download patch file Download diff file Expand all files Collapse all files

26
crates/ruff_db/src/source.rs
View file

@ -7,6 +7,7 @@ use ruff_source_file::LineIndex;
use crate::Db;
use crate::files::{File, FilePath};
use crate::system::System;
/// Reads the source text of a python text file (must be valid UTF8) or notebook.
#[salsa::tracked(heap_size=ruff_memory_usage::heap_size)]
@ -15,7 +16,7 @@ pub fn source_text(db: &dyn Db, file: File) -> SourceText {
let _span = tracing::trace_span!("source_text", file = %path).entered();
let mut read_error = None;
let kind = if is_notebook(file.path(db)) {
let kind = if is_notebook(db.system(), path) {
file.read_to_notebook(db)
.unwrap_or_else(|error| {
tracing::debug!("Failed to read notebook '{path}': {error}");
@ -40,18 +41,17 @@ pub fn source_text(db: &dyn Db, file: File) -> SourceText {
}
}
fn is_notebook(path: &FilePath) -> bool {
match path {
FilePath::System(system) => system.extension().is_some_and(|extension| {
PySourceType::try_from_extension(extension) == Some(PySourceType::Ipynb)
}),
FilePath::SystemVirtual(system_virtual) => {
system_virtual.extension().is_some_and(|extension| {
PySourceType::try_from_extension(extension) == Some(PySourceType::Ipynb)
})
}
FilePath::Vendored(_) => false,
}
fn is_notebook(system: &dyn System, path: &FilePath) -> bool {
let source_type = match path {
FilePath::System(path) => system.source_type(path),
FilePath::SystemVirtual(system_virtual) => system.virtual_path_source_type(system_virtual),
FilePath::Vendored(_) => return false,
};
let with_extension_fallback =
source_type.or_else(|| PySourceType::try_from_extension(path.extension()?));
with_extension_fallback == Some(PySourceType::Ipynb)
}
/// The source text of a file containing python code.

33
crates/ruff_db/src/system.rs
View file

@ -9,6 +9,7 @@ pub use os::OsSystem;
use filetime::FileTime;
use ruff_notebook::{Notebook, NotebookError};
use ruff_python_ast::PySourceType;
use std::error::Error;
use std::fmt::{Debug, Formatter};
use std::path::{Path, PathBuf};
@ -16,12 +17,11 @@ use std::{fmt, io};
pub use test::{DbWithTestSystem, DbWithWritableSystem, InMemorySystem, TestSystem};
use walk_directory::WalkDirectoryBuilder;
use crate::file_revision::FileRevision;
pub use self::path::{
DeduplicatedNestedPathsIter, SystemPath, SystemPathBuf, SystemVirtualPath,
SystemVirtualPathBuf, deduplicate_nested_paths,
};
use crate::file_revision::FileRevision;
mod memory_fs;
#[cfg(feature = "os")]
@ -66,6 +66,35 @@ pub trait System: Debug + Sync + Send {
/// See [dunce::canonicalize] for more information.
fn canonicalize_path(&self, path: &SystemPath) -> Result<SystemPathBuf>;
/// Returns the source type for `path` if known or `None`.
///
/// The default is to always return `None`, assuming the system
/// has no additional information and that the caller should
/// rely on the file extension instead.
///
/// This is primarily used for the LSP integration to respect
/// the chosen language (or the fact that it is a notebook) in
/// the editor.
fn source_type(&self, path: &SystemPath) -> Option<PySourceType> {
let _ = path;
None
}
/// Returns the source type for `path` if known or `None`.
///
/// The default is to always return `None`, assuming the system
/// has no additional information and that the caller should
/// rely on the file extension instead.
///
/// This is primarily used for the LSP integration to respect
/// the chosen language (or the fact that it is a notebook) in
/// the editor.
fn virtual_path_source_type(&self, path: &SystemVirtualPath) -> Option<PySourceType> {
let _ = path;
None
}
/// Reads the content of the file at `path` into a [`String`].
fn read_to_string(&self, path: &SystemPath) -> Result<String>;

6
crates/ruff_notebook/src/index.rs
View file

@ -39,7 +39,7 @@ impl NotebookIndex {
/// Returns an iterator over the starting rows of each cell (1-based).
///
/// This yields one entry per Python cell (skipping over Makrdown cell).
/// This yields one entry per Python cell (skipping over Markdown cell).
pub fn iter(&self) -> impl Iterator<Item = CellStart> + '_ {
self.cell_starts.iter().copied()
}
@ -47,7 +47,7 @@ impl NotebookIndex {
/// Translates the given [`LineColumn`] based on the indexing table.
///
/// This will translate the row/column in the concatenated source code
/// to the row/column in the Jupyter Notebook.
/// to the row/column in the Jupyter Notebook cell.
pub fn translate_line_column(&self, source_location: &LineColumn) -> LineColumn {
LineColumn {
line: self
@ -60,7 +60,7 @@ impl NotebookIndex {
/// Translates the given [`SourceLocation`] based on the indexing table.
///
/// This will translate the line/character in the concatenated source code
/// to the line/character in the Jupyter Notebook.
/// to the line/character in the Jupyter Notebook cell.
pub fn translate_source_location(&self, source_location: &SourceLocation) -> SourceLocation {
SourceLocation {
line: self

19
crates/ruff_notebook/src/notebook.rs
View file

@ -13,7 +13,7 @@ use thiserror::Error;
use ruff_diagnostics::{SourceMap, SourceMarker};
use ruff_source_file::{NewlineWithTrailingNewline, OneIndexed, UniversalNewlineIterator};
use ruff_text_size::TextSize;
use ruff_text_size::{TextRange, TextSize};
use crate::cell::CellOffsets;
use crate::index::NotebookIndex;
@ -294,7 +294,7 @@ impl Notebook {
}
}
/// Build and return the [`JupyterIndex`].
/// Build and return the [`NotebookIndex`].
///
/// ## Notes
///
@ -388,6 +388,21 @@ impl Notebook {
&self.cell_offsets
}
/// Returns the start offset of the cell at index `cell` in the concatenated
/// text document.
pub fn cell_offset(&self, cell: OneIndexed) -> Option<TextSize> {
self.cell_offsets.get(cell.to_zero_indexed()).copied()
}
/// Returns the text range in the concatenated document of the cell
/// with index `cell`.
pub fn cell_range(&self, cell: OneIndexed) -> Option<TextRange> {
let start = self.cell_offsets.get(cell.to_zero_indexed()).copied()?;
let end = self.cell_offsets.get(cell.to_zero_indexed() + 1).copied()?;
Some(TextRange::new(start, end))
}
/// Return `true` if the notebook has a trailing newline, `false` otherwise.
pub fn trailing_newline(&self) -> bool {
self.trailing_newline

18
crates/ty_server/src/capabilities.rs
View file

@ -1,10 +1,10 @@
use lsp_types::{
ClientCapabilities, CompletionOptions, DeclarationCapability, DiagnosticOptions,
DiagnosticServerCapabilities, HoverProviderCapability, InlayHintOptions,
InlayHintServerCapabilities, MarkupKind, OneOf, RenameOptions,
SelectionRangeProviderCapability, SemanticTokensFullOptions, SemanticTokensLegend,
SemanticTokensOptions, SemanticTokensServerCapabilities, ServerCapabilities,
SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
InlayHintServerCapabilities, MarkupKind, NotebookCellSelector, NotebookSelector, OneOf,
RenameOptions, SelectionRangeProviderCapability, SemanticTokensFullOptions,
SemanticTokensLegend, SemanticTokensOptions, SemanticTokensServerCapabilities,
ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
TextDocumentSyncOptions, TypeDefinitionProviderCapability, WorkDoneProgressOptions,
};
@ -422,6 +422,16 @@ pub(crate) fn server_capabilities(
selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
document_symbol_provider: Some(OneOf::Left(true)),
workspace_symbol_provider: Some(OneOf::Left(true)),
notebook_document_sync: Some(OneOf::Left(lsp_types::NotebookDocumentSyncOptions {
save: Some(false),
notebook_selector: [NotebookSelector::ByCells {
notebook: None,
cells: vec![NotebookCellSelector {
language: "python".to_string(),
}],
}]
.to_vec(),
})),
..Default::default()
}
}

20
crates/ty_server/src/document.rs
View file

@ -5,15 +5,15 @@ mod notebook;
mod range;
mod text_document;
pub(crate) use location::ToLink;
use lsp_types::{PositionEncodingKind, Url};
use ruff_db::system::{SystemPathBuf, SystemVirtualPath, SystemVirtualPathBuf};
use crate::system::AnySystemPath;
pub(crate) use location::ToLink;
pub use notebook::NotebookDocument;
pub(crate) use range::{FileRangeExt, PositionExt, RangeExt, TextSizeExt, ToRangeExt};
use ruff_db::system::{SystemPathBuf, SystemVirtualPath};
pub(crate) use text_document::DocumentVersion;
pub use text_document::TextDocument;
pub(crate) use text_document::{DocumentVersion, LanguageId};
/// A convenient enumeration for supported text encodings. Can be converted to [`lsp_types::PositionEncodingKind`].
// Please maintain the order from least to greatest priority for the derived `Ord` impl.
@ -84,13 +84,6 @@ impl DocumentKey {
}
}
pub(crate) fn as_opaque(&self) -> Option<&str> {
match self {
Self::Opaque(uri) => Some(uri),
Self::File(_) => None,
}
}
/// Returns the corresponding [`AnySystemPath`] for this document key.
///
/// Note, calling this method on a `DocumentKey::Opaque` representing a cell document
@ -104,6 +97,13 @@ impl DocumentKey {
}
}
}
pub(super) fn into_file_path(self) -> AnySystemPath {
match self {
Self::File(path) => AnySystemPath::System(path),
Self::Opaque(uri) => AnySystemPath::SystemVirtual(SystemVirtualPathBuf::from(uri)),
}
}
}
impl From<AnySystemPath> for DocumentKey {

349
crates/ty_server/src/document/notebook.rs
View file

@ -1,26 +1,29 @@
use anyhow::Ok;
use lsp_types::NotebookCellKind;
use ruff_notebook::CellMetadata;
use rustc_hash::{FxBuildHasher, FxHashMap};
use ruff_source_file::OneIndexed;
use rustc_hash::FxHashMap;
use super::DocumentVersion;
use crate::{PositionEncoding, TextDocument};
use super::{DocumentKey, DocumentVersion};
use crate::session::index::Index;
pub(super) type CellId = usize;
/// The state of a notebook document in the server. Contains an array of cells whose
/// contents are internally represented by [`TextDocument`]s.
/// A notebook document.
///
/// This notebook document only stores the metadata about the notebook
/// and the cell metadata. The cell contents are stored as separate
/// [`super::TextDocument`]s (they can be looked up by the Cell's URL).
#[derive(Clone, Debug)]
pub struct NotebookDocument {
url: lsp_types::Url,
cells: Vec<NotebookCell>,
metadata: ruff_notebook::RawNotebookMetadata,
version: DocumentVersion,
// Used to quickly find the index of a cell for a given URL.
cell_index: FxHashMap<String, CellId>,
/// Map from Cell URL to their index in `cells`
cell_index: FxHashMap<lsp_types::Url, usize>,
}
/// A single cell within a notebook, which has text contents represented as a `TextDocument`.
/// The metadata of a single cell within a notebook.
///
/// The cell's content are stored as a [`TextDocument`] and can be looked up by the Cell's URL.
#[derive(Clone, Debug)]
struct NotebookCell {
/// The URL uniquely identifying the cell.
@ -33,7 +36,7 @@ struct NotebookCell {
/// > <https://microsoft.github.io/language-server-protocol/specifications/lsp/3.18/specification/#notebookDocument_synchronization>
url: lsp_types::Url,
kind: NotebookCellKind,
document: TextDocument,
execution_summary: Option<lsp_types::ExecutionSummary>,
}
impl NotebookDocument {
@ -42,32 +45,18 @@ impl NotebookDocument {
notebook_version: DocumentVersion,
cells: Vec<lsp_types::NotebookCell>,
metadata: serde_json::Map<String, serde_json::Value>,
cell_documents: Vec<lsp_types::TextDocumentItem>,
) -> crate::Result<Self> {
let mut cells: Vec<_> = cells.into_iter().map(NotebookCell::empty).collect();
let cell_index = Self::make_cell_index(&cells);
for cell_document in cell_documents {
let index = cell_index
.get(cell_document.uri.as_str())
.copied()
.ok_or_else(|| {
anyhow::anyhow!(
"Received content for cell `{}` that isn't present in the metadata",
cell_document.uri
)
})?;
cells[index].document =
TextDocument::new(cell_document.uri, cell_document.text, cell_document.version)
.with_language_id(&cell_document.language_id);
}
let cells: Vec<_> = cells.into_iter().map(NotebookCell::new).collect();
let index = cells
.iter()
.enumerate()
.map(|(index, cell)| (cell.url.clone(), index))
.collect();
Ok(Self {
cell_index: index,
url,
version: notebook_version,
cell_index,
cells,
metadata: serde_json::from_value(serde_json::Value::Object(metadata))?,
})
@ -79,30 +68,46 @@ impl NotebookDocument {
/// Generates a pseudo-representation of a notebook that lacks per-cell metadata and contextual information
/// but should still work with Ruff's linter.
pub fn make_ruff_notebook(&self) -> ruff_notebook::Notebook {
pub(crate) fn to_ruff_notebook(&self, index: &Index) -> ruff_notebook::Notebook {
let cells = self
.cells
.iter()
.map(|cell| match cell.kind {
.map(|cell| {
let cell_text =
if let Ok(document) = index.document(&DocumentKey::from_url(&cell.url)) {
if let Some(text_document) = document.as_text() {
Some(text_document.contents().to_string())
} else {
tracing::warn!("Non-text document found for cell `{}`", cell.url);
None
}
} else {
tracing::warn!("Text document not found for cell `{}`", cell.url);
None
}
.unwrap_or_default();
let source = ruff_notebook::SourceValue::String(cell_text);
match cell.kind {
NotebookCellKind::Code => ruff_notebook::Cell::Code(ruff_notebook::CodeCell {
execution_count: None,
execution_count: cell
.execution_summary
.as_ref()
.map(|summary| i64::from(summary.execution_order)),
id: None,
metadata: CellMetadata::default(),
outputs: vec![],
source: ruff_notebook::SourceValue::String(
cell.document.contents().to_string(),
),
source,
}),
NotebookCellKind::Markup => {
ruff_notebook::Cell::Markdown(ruff_notebook::MarkdownCell {
attachments: None,
id: None,
metadata: CellMetadata::default(),
source: ruff_notebook::SourceValue::String(
cell.document.contents().to_string(),
),
source,
})
}
}
})
.collect();
let raw_notebook = ruff_notebook::RawNotebook {
@ -118,93 +123,38 @@ impl NotebookDocument {
pub(crate) fn update(
&mut self,
cells: Option<lsp_types::NotebookDocumentCellChange>,
array: lsp_types::NotebookCellArrayChange,
updated_cells: Vec<lsp_types::NotebookCell>,
metadata_change: Option<serde_json::Map<String, serde_json::Value>>,
version: DocumentVersion,
encoding: PositionEncoding,
) -> crate::Result<()> {
self.version = version;
if let Some(lsp_types::NotebookDocumentCellChange {
structure,
data,
text_content,
}) = cells
{
// The structural changes should be done first, as they may affect the cell index.
if let Some(structure) = structure {
let start = structure.array.start as usize;
let delete = structure.array.delete_count as usize;
let new_cells = array.cells.unwrap_or_default();
let start = array.start as usize;
// This is required because of the way the `NotebookCell` is modelled. We include
// the `TextDocument` within the `NotebookCell` so when it's deleted, the
// corresponding `TextDocument` is removed as well. But, when cells are
// re-ordered, the change request doesn't provide the actual contents of the cell.
// Instead, it only provides that (a) these cell URIs were removed, and (b) these
// cell URIs were added.
// https://github.com/astral-sh/ruff/issues/12573
let mut deleted_cells = FxHashMap::default();
let added = new_cells.len();
let deleted_range = start..start + array.delete_count as usize;
// First, delete the cells and remove them from the index.
if delete > 0 {
for cell in self.cells.drain(start..start + delete) {
self.cell_index.remove(cell.url.as_str());
deleted_cells.insert(cell.url, cell.document);
}
}
self.cells.splice(
deleted_range.clone(),
new_cells.into_iter().map(NotebookCell::new),
);
// Second, insert the new cells with the available information. This array does not
// provide the actual contents of the cells, so we'll initialize them with empty
// contents.
for cell in structure.array.cells.into_iter().flatten().rev() {
let (content, version) =
if let Some(text_document) = deleted_cells.remove(&cell.document) {
let version = text_document.version();
(text_document.into_contents(), version)
} else {
(String::new(), 0)
};
// Re-build the cell-index if new cells were added, deleted or removed
if !deleted_range.is_empty() || added > 0 {
self.cell_index.clear();
self.cell_index.extend(
self.cells
.insert(start, NotebookCell::new(cell, content, version));
}
// Third, register the new cells in the index and update existing ones that came
// after the insertion.
for (index, cell) in self.cells.iter().enumerate().skip(start) {
self.cell_index.insert(cell.url.to_string(), index);
}
// Finally, update the text document that represents the cell with the actual
// contents. This should be done at the end so that both the `cells` and
// `cell_index` are updated before we start applying the changes to the cells.
if let Some(did_open) = structure.did_open {
for cell_text_document in did_open {
if let Some(cell) = self.cell_by_uri_mut(cell_text_document.uri.as_str()) {
cell.document = TextDocument::new(
cell_text_document.uri,
cell_text_document.text,
cell_text_document.version,
.iter()
.enumerate()
.map(|(i, cell)| (cell.url.clone(), i)),
);
}
}
}
}
if let Some(cell_data) = data {
for cell in cell_data {
if let Some(existing_cell) = self.cell_by_uri_mut(cell.document.as_str()) {
existing_cell.kind = cell.kind;
}
}
}
if let Some(content_changes) = text_content {
for content_change in content_changes {
if let Some(cell) = self.cell_by_uri_mut(content_change.document.uri.as_str()) {
cell.document
.apply_changes(content_change.changes, version, encoding);
}
}
for cell in updated_cells {
if let Some(existing_cell_index) = self.cell_index.get(&cell.document).copied() {
self.cells[existing_cell_index].kind = cell.kind;
}
}
@ -221,16 +171,10 @@ impl NotebookDocument {
}
/// Get the URI for a cell by its index within the cell array.
pub(crate) fn cell_uri_by_index(&self, index: CellId) -> Option<&lsp_types::Url> {
self.cells.get(index).map(|cell| &cell.url)
}
/// Get the text document representing the contents of a cell by the cell URI.
#[expect(unused)]
pub(crate) fn cell_document_by_uri(&self, uri: &str) -> Option<&TextDocument> {
pub(crate) fn cell_uri_by_index(&self, index: OneIndexed) -> Option<&lsp_types::Url> {
self.cells
.get(*self.cell_index.get(uri)?)
.map(|cell| &cell.document)
.get(index.to_zero_indexed())
.map(|cell| &cell.url)
}
/// Returns a list of cell URIs in the order they appear in the array.
@ -238,160 +182,19 @@ impl NotebookDocument {
self.cells.iter().map(|cell| &cell.url)
}
fn cell_by_uri_mut(&mut self, uri: &str) -> Option<&mut NotebookCell> {
self.cells.get_mut(*self.cell_index.get(uri)?)
}
fn make_cell_index(cells: &[NotebookCell]) -> FxHashMap<String, CellId> {
let mut index = FxHashMap::with_capacity_and_hasher(cells.len(), FxBuildHasher);
for (i, cell) in cells.iter().enumerate() {
index.insert(cell.url.to_string(), i);
}
index
pub(crate) fn cell_index_by_uri(&self, cell_url: &lsp_types::Url) -> Option<OneIndexed> {
Some(OneIndexed::from_zero_indexed(
self.cell_index.get(cell_url).copied()?,
))
}
}
impl NotebookCell {
pub(crate) fn empty(cell: lsp_types::NotebookCell) -> Self {
pub(crate) fn new(cell: lsp_types::NotebookCell) -> Self {
Self {
kind: cell.kind,
document: TextDocument::new(
cell.document.clone(),
String::new(),
DocumentVersion::default(),
),
url: cell.document,
}
}
pub(crate) fn new(
cell: lsp_types::NotebookCell,
contents: String,
version: DocumentVersion,
) -> Self {
Self {
document: TextDocument::new(cell.document.clone(), contents, version),
url: cell.document,
kind: cell.kind,
execution_summary: cell.execution_summary,
}
}
}
#[cfg(test)]
mod tests {
use super::NotebookDocument;
enum TestCellContent {
#[expect(dead_code)]
Markup(String),
Code(String),
}
fn create_test_url(index: usize) -> lsp_types::Url {
lsp_types::Url::parse(&format!("cell:/test.ipynb#{index}")).unwrap()
}
fn create_test_notebook(test_cells: Vec<TestCellContent>) -> NotebookDocument {
let mut cells = Vec::with_capacity(test_cells.len());
let mut cell_documents = Vec::with_capacity(test_cells.len());
for (index, test_cell) in test_cells.into_iter().enumerate() {
let url = create_test_url(index);
match test_cell {
TestCellContent::Markup(content) => {
cells.push(lsp_types::NotebookCell {
kind: lsp_types::NotebookCellKind::Markup,
document: url.clone(),
metadata: None,
execution_summary: None,
});
cell_documents.push(lsp_types::TextDocumentItem {
uri: url,
language_id: "markdown".to_owned(),
version: 0,
text: content,
});
}
TestCellContent::Code(content) => {
cells.push(lsp_types::NotebookCell {
kind: lsp_types::NotebookCellKind::Code,
document: url.clone(),
metadata: None,
execution_summary: None,
});
cell_documents.push(lsp_types::TextDocumentItem {
uri: url,
language_id: "python".to_owned(),
version: 0,
text: content,
});
}
}
}
NotebookDocument::new(
lsp_types::Url::parse("file://test.ipynb").unwrap(),
0,
cells,
serde_json::Map::default(),
cell_documents,
)
.unwrap()
}
/// This test case checks that for a notebook with three code cells, when the client sends a
/// change request to swap the first two cells, the notebook document is updated correctly.
///
/// The swap operation as a change request is represented as deleting the first two cells and
/// adding them back in the reverse order.
#[test]
fn swap_cells() {
let mut notebook = create_test_notebook(vec![
TestCellContent::Code("cell = 0".to_owned()),
TestCellContent::Code("cell = 1".to_owned()),
TestCellContent::Code("cell = 2".to_owned()),
]);
notebook
.update(
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: lsp_types::NotebookCellKind::Code,
document: create_test_url(1),
metadata: None,
execution_summary: None,
},
lsp_types::NotebookCell {
kind: lsp_types::NotebookCellKind::Code,
document: create_test_url(0),
metadata: None,
execution_summary: None,
},
]),
},
did_open: None,
did_close: None,
}),
data: None,
text_content: None,
}),
None,
1,
crate::PositionEncoding::default(),
)
.unwrap();
assert_eq!(
notebook.make_ruff_notebook().source_code(),
"cell = 1
cell = 0
cell = 2
"
);
}
}

72
crates/ty_server/src/document/range.rs
View file

@ -78,7 +78,7 @@ impl LspPosition {
}
pub(crate) trait RangeExt {
/// Convert an LSP Range to internal [`TextRange`].
/// Convert an LSP Range to a [`TextRange`].
///
/// Returns `None` if `file` is a notebook and the
/// cell identified by `url` can't be looked up or if the notebook
@ -110,6 +110,10 @@ impl RangeExt for lsp_types::Range {
pub(crate) trait PositionExt {
/// Convert an LSP Position to internal `TextSize`.
///
/// For notebook support, this uses the URI to determine which cell the position
/// refers to, and maps the cell-relative position to the absolute position in the
/// concatenated notebook file.
///
/// Returns `None` if `file` is a notebook and the
/// cell identified by `url` can't be looked up or if the notebook
/// isn't open in the editor.
@ -127,18 +131,39 @@ impl PositionExt for lsp_types::Position {
&self,
db: &dyn Db,
file: File,
_url: &lsp_types::Url,
url: &lsp_types::Url,
encoding: PositionEncoding,
) -> Option<TextSize> {
let source = source_text(db, file);
let index = line_index(db, file);
if let Some(notebook) = source.as_notebook() {
let notebook_document = db.notebook_document(file)?;
let cell_index = notebook_document.cell_index_by_uri(url)?;
let cell_start_offset = notebook.cell_offset(cell_index).unwrap_or_default();
let cell_relative_line = OneIndexed::from_zero_indexed(u32_index_to_usize(self.line));
let cell_start_location =
index.source_location(cell_start_offset, source.as_str(), encoding.into());
assert_eq!(cell_start_location.character_offset, OneIndexed::MIN);
// Absolute position into the concatenated notebook source text.
let absolute_position = SourceLocation {
line: cell_start_location
.line
.saturating_add(cell_relative_line.to_zero_indexed()),
character_offset: OneIndexed::from_zero_indexed(u32_index_to_usize(self.character)),
};
return Some(index.offset(absolute_position, &source, encoding.into()));
}
Some(lsp_position_to_text_size(*self, &source, &index, encoding))
}
}
pub(crate) trait TextSizeExt {
/// Converts self into a position into an LSP text document (can be a cell or regular document).
/// Converts `self` into a position in an LSP text document (can be a cell or regular document).
///
/// Returns `None` if the position can't be converted:
///
@ -165,6 +190,19 @@ impl TextSizeExt for TextSize {
let source = source_text(db, file);
let index = line_index(db, file);
if let Some(notebook) = source.as_notebook() {
let notebook_document = db.notebook_document(file)?;
let start = index.source_location(*self, source.as_str(), encoding.into());
let cell = notebook.index().cell(start.line)?;
let cell_relative_start = notebook.index().translate_source_location(&start);
return Some(LspPosition {
uri: Some(notebook_document.cell_uri_by_index(cell)?.clone()),
position: source_location_to_position(&cell_relative_start),
});
}
let uri = file_to_url(db, file);
let position = text_size_to_lsp_position(*self, &source, &index, encoding);
@ -252,6 +290,34 @@ impl ToRangeExt for TextRange {
) -> Option<LspRange> {
let source = source_text(db, file);
let index = line_index(db, file);
if let Some(notebook) = source.as_notebook() {
let notebook_index = notebook.index();
let notebook_document = db.notebook_document(file)?;
let start_in_concatenated =
index.source_location(self.start(), &source, encoding.into());
let cell_index = notebook_index.cell(start_in_concatenated.line)?;
let end_in_concatenated = index.source_location(self.end(), &source, encoding.into());
let start_in_cell = source_location_to_position(
&notebook_index.translate_source_location(&start_in_concatenated),
);
let end_in_cell = source_location_to_position(
&notebook_index.translate_source_location(&end_in_concatenated),
);
let cell_uri = notebook_document
.cell_uri_by_index(cell_index)
.expect("Index to contain an URI for every cell");
return Some(LspRange {
uri: Some(cell_uri.clone()),
range: lsp_types::Range::new(start_in_cell, end_in_cell),
});
}
let range = text_range_to_lsp_range(*self, &source, &index, encoding);
let uri = file_to_url(db, file);

54
crates/ty_server/src/document/text_document.rs
View file

@ -2,11 +2,13 @@ use lsp_types::{TextDocumentContentChangeEvent, Url};
use ruff_source_file::LineIndex;
use crate::PositionEncoding;
use super::range::lsp_range_to_text_range;
use crate::document::range::lsp_range_to_text_range;
use crate::system::AnySystemPath;
pub(crate) type DocumentVersion = i32;
/// A regular text file or the content of a notebook cell.
///
/// The state of an individual document in the server. Stays up-to-date
/// with changes made by the user, including unsaved changes.
#[derive(Debug, Clone)]
@ -16,15 +18,16 @@ pub struct TextDocument {
/// The string contents of the document.
contents: String,
/// A computed line index for the document. This should always reflect
/// the current version of `contents`. Using a function like [`Self::modify`]
/// will re-calculate the line index automatically when the `contents` value is updated.
index: LineIndex,
/// The latest version of the document, set by the LSP client. The server will panic in
/// debug mode if we attempt to update the document with an 'older' version.
version: DocumentVersion,
/// The language ID of the document as provided by the client.
language_id: Option<LanguageId>,
/// For cells, the path to the notebook document.
notebook: Option<AnySystemPath>,
}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
@ -44,13 +47,12 @@ impl From<&str> for LanguageId {
impl TextDocument {
pub fn new(url: Url, contents: String, version: DocumentVersion) -> Self {
let index = LineIndex::from_source_text(&contents);
Self {
url,
contents,
index,
version,
language_id: None,
notebook: None,
}
}
@ -60,6 +62,12 @@ impl TextDocument {
self
}
#[must_use]
pub(crate) fn with_notebook(mut self, notebook: AnySystemPath) -> Self {
self.notebook = Some(notebook);
self
}
pub fn into_contents(self) -> String {
self.contents
}
@ -72,10 +80,6 @@ impl TextDocument {
&self.contents
}
pub fn index(&self) -> &LineIndex {
&self.index
}
pub fn version(&self) -> DocumentVersion {
self.version
}
@ -84,6 +88,10 @@ impl TextDocument {
self.language_id
}
pub(crate) fn notebook(&self) -> Option<&AnySystemPath> {
self.notebook.as_ref()
}
pub fn apply_changes(
&mut self,
changes: Vec<lsp_types::TextDocumentContentChangeEvent>,
@ -105,7 +113,7 @@ impl TextDocument {
}
let mut new_contents = self.contents().to_string();
let mut active_index = self.index().clone();
let mut active_index = LineIndex::from_source_text(&new_contents);
for TextDocumentContentChangeEvent {
range,
@ -127,34 +135,22 @@ impl TextDocument {
active_index = LineIndex::from_source_text(&new_contents);
}
self.modify_with_manual_index(|contents, version, index| {
*index = active_index;
self.modify(|contents, version| {
*contents = new_contents;
*version = new_version;
});
}
pub fn update_version(&mut self, new_version: DocumentVersion) {
self.modify_with_manual_index(|_, version, _| {
self.modify(|_, version| {
*version = new_version;
});
}
// A private function for modifying the document's internal state
fn modify(&mut self, func: impl FnOnce(&mut String, &mut DocumentVersion)) {
self.modify_with_manual_index(|c, v, i| {
func(c, v);
*i = LineIndex::from_source_text(c);
});
}
// A private function for overriding how we update the line index by default.
fn modify_with_manual_index(
&mut self,
func: impl FnOnce(&mut String, &mut DocumentVersion, &mut LineIndex),
) {
fn modify(&mut self, func: impl FnOnce(&mut String, &mut DocumentVersion)) {
let old_version = self.version;
func(&mut self.contents, &mut self.version, &mut self.index);
func(&mut self.contents, &mut self.version);
debug_assert!(self.version >= old_version);
}
}

7
crates/ty_server/src/server/api.rs
View file

@ -147,6 +147,9 @@ pub(super) fn notification(notif: server::Notification) -> Task {
notifications::DidOpenNotebookHandler::METHOD => {
sync_notification_task::<notifications::DidOpenNotebookHandler>(notif)
}
notifications::DidChangeNotebookHandler::METHOD => {
sync_notification_task::<notifications::DidChangeNotebookHandler>(notif)
}
notifications::DidCloseNotebookHandler::METHOD => {
sync_notification_task::<notifications::DidCloseNotebookHandler>(notif)
}
@ -273,8 +276,8 @@ where
});
};
let path = document.to_file_path();
let db = session.project_db(&path).clone();
let path = document.notebook_or_file_path();
let db = session.project_db(path).clone();
Box::new(move |client| {
let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();

131
crates/ty_server/src/server/api/diagnostics.rs
View file

@ -3,29 +3,30 @@ use std::hash::{DefaultHasher, Hash as _, Hasher as _};
use lsp_types::notification::PublishDiagnostics;
use lsp_types::{
CodeDescription, Diagnostic, DiagnosticRelatedInformation, DiagnosticSeverity, DiagnosticTag,
NumberOrString, PublishDiagnosticsParams, Range, Url,
NumberOrString, PublishDiagnosticsParams, Url,
};
use ruff_db::source::source_text;
use rustc_hash::FxHashMap;
use ruff_db::diagnostic::{Annotation, Severity, SubDiagnostic};
use ruff_db::files::FileRange;
use ruff_db::files::{File, FileRange};
use ruff_db::system::SystemPathBuf;
use ty_project::{Db as _, ProjectDatabase};
use crate::Db;
use crate::document::{FileRangeExt, ToRangeExt};
use crate::session::DocumentSnapshot;
use crate::session::DocumentHandle;
use crate::session::client::Client;
use crate::system::{AnySystemPath, file_to_url};
use crate::{NotebookDocument, PositionEncoding, Session};
use crate::{PositionEncoding, Session};
pub(super) struct Diagnostics<'a> {
pub(super) struct Diagnostics {
items: Vec<ruff_db::diagnostic::Diagnostic>,
encoding: PositionEncoding,
notebook: Option<&'a NotebookDocument>,
file_or_notebook: File,
}
impl Diagnostics<'_> {
impl Diagnostics {
/// Computes the result ID for `diagnostics`.
///
/// Returns `None` if there are no diagnostics.
@ -53,30 +54,27 @@ impl Diagnostics<'_> {
}
pub(super) fn to_lsp_diagnostics(&self, db: &ProjectDatabase) -> LspDiagnostics {
if let Some(notebook) = self.notebook {
if let Some(notebook_document) = db.notebook_document(self.file_or_notebook) {
let mut cell_diagnostics: FxHashMap<Url, Vec<Diagnostic>> = FxHashMap::default();
// Populates all relevant URLs with an empty diagnostic list. This ensures that documents
// without diagnostics still get updated.
for cell_url in notebook.cell_urls() {
for cell_url in notebook_document.cell_urls() {
cell_diagnostics.entry(cell_url.clone()).or_default();
}
for (cell_index, diagnostic) in self.items.iter().map(|diagnostic| {
(
// TODO: Use the cell index instead using `SourceKind`
usize::default(),
to_lsp_diagnostic(db, diagnostic, self.encoding),
)
}) {
let Some(cell_uri) = notebook.cell_uri_by_index(cell_index) else {
tracing::warn!("Unable to find notebook cell at index {cell_index}");
for diagnostic in &self.items {
let (url, lsp_diagnostic) = to_lsp_diagnostic(db, diagnostic, self.encoding);
let Some(url) = url else {
tracing::warn!("Unable to find notebook cell");
continue;
};
cell_diagnostics
.entry(cell_uri.clone())
.entry(url)
.or_default()
.push(diagnostic);
.push(lsp_diagnostic);
}
LspDiagnostics::NotebookDocument(cell_diagnostics)
@ -84,7 +82,7 @@ impl Diagnostics<'_> {
LspDiagnostics::TextDocument(
self.items
.iter()
.map(|diagnostic| to_lsp_diagnostic(db, diagnostic, self.encoding))
.map(|diagnostic| to_lsp_diagnostic(db, diagnostic, self.encoding).1)
.collect(),
)
}
@ -115,16 +113,25 @@ impl LspDiagnostics {
}
}
pub(super) fn clear_diagnostics_if_needed(
document: &DocumentHandle,
session: &Session,
client: &Client,
) {
if session.client_capabilities().supports_pull_diagnostics() && !document.is_cell_or_notebook()
{
return;
}
clear_diagnostics(document.url(), client);
}
/// Clears the diagnostics for the document identified by `uri`.
///
/// This is done by notifying the client with an empty list of diagnostics for the document.
/// For notebook cells, this clears diagnostics for the specific cell.
/// For other document types, this clears diagnostics for the main document.
pub(super) fn clear_diagnostics(session: &Session, uri: &lsp_types::Url, client: &Client) {
if session.client_capabilities().supports_pull_diagnostics() {
return;
}
pub(super) fn clear_diagnostics(uri: &lsp_types::Url, client: &Client) {
client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
uri: uri.clone(),
diagnostics: vec![],
@ -135,25 +142,30 @@ pub(super) fn clear_diagnostics(session: &Session, uri: &lsp_types::Url, client:
/// Publishes the diagnostics for the given document snapshot using the [publish diagnostics
/// notification] .
///
/// This function is a no-op if the client supports pull diagnostics.
/// Unlike [`publish_diagnostics`], this function only publishes diagnostics if a client doesn't support
/// pull diagnostics and `document` is not a notebook or cell (VS Code
/// does not support pull diagnostics for notebooks or cells (as of 2025-11-12).
///
/// [publish diagnostics notification]: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_publishDiagnostics
pub(super) fn publish_diagnostics(session: &Session, url: &lsp_types::Url, client: &Client) {
if session.client_capabilities().supports_pull_diagnostics() {
pub(super) fn publish_diagnostics_if_needed(
document: &DocumentHandle,
session: &Session,
client: &Client,
) {
if !document.is_cell_or_notebook() && session.client_capabilities().supports_pull_diagnostics()
{
return;
}
let snapshot = match session.snapshot_document(url) {
Ok(document) => document,
Err(err) => {
tracing::debug!("Failed to resolve document for URL `{}`: {}", url, err);
return;
publish_diagnostics(document, session, client);
}
};
let db = session.project_db(&snapshot.to_file_path());
/// Publishes the diagnostics for the given document snapshot using the [publish diagnostics
/// notification].
pub(super) fn publish_diagnostics(document: &DocumentHandle, session: &Session, client: &Client) {
let db = session.project_db(document.notebook_or_file_path());
let Some(diagnostics) = compute_diagnostics(db, &snapshot) else {
let Some(diagnostics) = compute_diagnostics(db, document, session.position_encoding()) else {
return;
};
@ -162,13 +174,13 @@ pub(super) fn publish_diagnostics(session: &Session, url: &lsp_types::Url, clien
client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
uri,
diagnostics,
version: Some(snapshot.document().version()),
version: Some(document.version()),
});
};
match diagnostics.to_lsp_diagnostics(db) {
LspDiagnostics::TextDocument(diagnostics) => {
publish_diagnostics_notification(url.clone(), diagnostics);
publish_diagnostics_notification(document.url().clone(), diagnostics);
}
LspDiagnostics::NotebookDocument(cell_diagnostics) => {
for (cell_url, diagnostics) in cell_diagnostics {
@ -238,7 +250,7 @@ pub(crate) fn publish_settings_diagnostics(
// Convert diagnostics to LSP format
let lsp_diagnostics = file_diagnostics
.into_iter()
.map(|diagnostic| to_lsp_diagnostic(db, &diagnostic, session_encoding))
.map(|diagnostic| to_lsp_diagnostic(db, &diagnostic, session_encoding).1)
.collect::<Vec<_>>();
client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
@ -249,24 +261,26 @@ pub(crate) fn publish_settings_diagnostics(
}
}
pub(super) fn compute_diagnostics<'a>(
pub(super) fn compute_diagnostics(
db: &ProjectDatabase,
snapshot: &'a DocumentSnapshot,
) -> Option<Diagnostics<'a>> {
let Some(file) = snapshot.to_file(db) else {
document: &DocumentHandle,
encoding: PositionEncoding,
) -> Option<Diagnostics> {
let Some(file) = document.notebook_or_file(db) else {
tracing::info!(
"No file found for snapshot for `{}`",
snapshot.to_file_path()
document.notebook_or_file_path()
);
return None;
};
tracing::debug!("source text: {}", source_text(db, file).as_str());
let diagnostics = db.check_file(file);
Some(Diagnostics {
items: diagnostics,
encoding: snapshot.encoding(),
notebook: snapshot.notebook(),
encoding,
file_or_notebook: file,
})
}
@ -276,16 +290,18 @@ pub(super) fn to_lsp_diagnostic(
db: &dyn Db,
diagnostic: &ruff_db::diagnostic::Diagnostic,
encoding: PositionEncoding,
) -> Diagnostic {
let range = if let Some(span) = diagnostic.primary_span() {
) -> (Option<lsp_types::Url>, Diagnostic) {
let location = diagnostic.primary_span().and_then(|span| {
let file = span.expect_ty_file();
span.range()
.and_then(|range| range.to_lsp_range(db, file, encoding))
span.range()?
.to_lsp_range(db, file, encoding)
.unwrap_or_default()
.local_range()
} else {
Range::default()
.to_location()
});
let (range, url) = match location {
Some(location) => (location.range, Some(location.uri)),
None => (lsp_types::Range::default(), None),
};
let severity = match diagnostic.severity() {
@ -341,6 +357,8 @@ pub(super) fn to_lsp_diagnostic(
);
}
(
url,
Diagnostic {
range,
severity: Some(severity),
@ -351,7 +369,8 @@ pub(super) fn to_lsp_diagnostic(
message: diagnostic.concise_message().to_string(),
related_information: Some(related_information),
data: None,
}
},
)
}
/// Converts an [`Annotation`] to a [`DiagnosticRelatedInformation`].

2
crates/ty_server/src/server/api/notifications.rs
View file

@ -1,5 +1,6 @@
mod cancel;
mod did_change;
mod did_change_notebook;
mod did_change_watched_files;
mod did_close;
mod did_close_notebook;
@ -8,6 +9,7 @@ mod did_open_notebook;
pub(super) use cancel::CancelNotificationHandler;
pub(super) use did_change::DidChangeTextDocumentHandler;
pub(super) use did_change_notebook::DidChangeNotebookHandler;
pub(super) use did_change_watched_files::DidChangeWatchedFiles;
pub(super) use did_close::DidCloseTextDocumentHandler;
pub(super) use did_close_notebook::DidCloseNotebookHandler;

18
crates/ty_server/src/server/api/notifications/did_change.rs
View file

@ -4,12 +4,10 @@ use lsp_types::{DidChangeTextDocumentParams, VersionedTextDocumentIdentifier};
use crate::server::Result;
use crate::server::api::LSPResult;
use crate::server::api::diagnostics::publish_diagnostics;
use crate::server::api::diagnostics::publish_diagnostics_if_needed;
use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
use crate::session::Session;
use crate::session::client::Client;
use crate::system::AnySystemPath;
use ty_project::watch::ChangeEvent;
pub(crate) struct DidChangeTextDocumentHandler;
@ -36,19 +34,7 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
.update_text_document(session, content_changes, version)
.with_failure_code(ErrorCode::InternalError)?;
let path = document.to_file_path();
let changes = match &*path {
AnySystemPath::System(system_path) => {
vec![ChangeEvent::file_content_changed(system_path.clone())]
}
AnySystemPath::SystemVirtual(virtual_path) => {
vec![ChangeEvent::ChangedVirtual(virtual_path.clone())]
}
};
session.apply_changes(&path, changes);
publish_diagnostics(session, document.url(), client);
publish_diagnostics_if_needed(&document, session, client);
Ok(())
}

40
crates/ty_server/src/server/api/notifications/did_change_notebook.rs Normal file
View file

@ -0,0 +1,40 @@
use lsp_server::ErrorCode;
use lsp_types as types;
use lsp_types::notification as notif;
use crate::server::Result;
use crate::server::api::LSPResult;
use crate::server::api::diagnostics::publish_diagnostics;
use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
use crate::session::Session;
use crate::session::client::Client;
pub(crate) struct DidChangeNotebookHandler;
impl NotificationHandler for DidChangeNotebookHandler {
type NotificationType = notif::DidChangeNotebookDocument;
}
impl SyncNotificationHandler for DidChangeNotebookHandler {
fn run(
session: &mut Session,
client: &Client,
types::DidChangeNotebookDocumentParams {
notebook_document: types::VersionedNotebookDocumentIdentifier { uri, version },
change: types::NotebookDocumentChangeEvent { cells, metadata },
}: types::DidChangeNotebookDocumentParams,
) -> Result<()> {
let document = session
.document_handle(&uri)
.with_failure_code(ErrorCode::InternalError)?;
document
.update_notebook_document(session, cells, metadata, version)
.with_failure_code(ErrorCode::InternalError)?;
// Always publish diagnostics because notebooks only support publish diagnostics.
publish_diagnostics(&document, session, client);
Ok(())
}
}

6
crates/ty_server/src/server/api/notifications/did_change_watched_files.rs
View file

@ -26,8 +26,7 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
let mut events_by_db: FxHashMap<_, Vec<ChangeEvent>> = FxHashMap::default();
for change in params.changes {
let key = DocumentKey::from_url(&change.uri);
let path = key.to_file_path();
let path = DocumentKey::from_url(&change.uri).into_file_path();
let system_path = match path {
AnySystemPath::System(system) => system,
@ -93,10 +92,9 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
);
} else {
for key in session.text_document_handles() {
publish_diagnostics(session, key.url(), client);
publish_diagnostics(&key, session, client);
}
}
// TODO: always publish diagnostics for notebook files (since they don't use pull diagnostics)
if client_capabilities.supports_inlay_hint_refresh() {
client.send_request::<types::request::InlayHintRefreshRequest>(session, (), |_, ()| {});

63
crates/ty_server/src/server/api/notifications/did_close.rs
View file

@ -1,15 +1,13 @@
use crate::server::Result;
use crate::server::api::LSPResult;
use crate::server::api::diagnostics::clear_diagnostics;
use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
use crate::session::Session;
use crate::session::client::Client;
use crate::system::AnySystemPath;
use lsp_server::ErrorCode;
use lsp_types::notification::DidCloseTextDocument;
use lsp_types::{DidCloseTextDocumentParams, TextDocumentIdentifier};
use ruff_db::Db as _;
use ty_project::Db as _;
use crate::server::Result;
use crate::server::api::LSPResult;
use crate::server::api::diagnostics::clear_diagnostics_if_needed;
use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
use crate::session::Session;
use crate::session::client::Client;
pub(crate) struct DidCloseTextDocumentHandler;
@ -31,53 +29,12 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
.document_handle(&uri)
.with_failure_code(ErrorCode::InternalError)?;
let path = document.to_file_path().into_owned();
let url = document.url().clone();
document
let should_clear_diagnostics = document
.close(session)
.with_failure_code(ErrorCode::InternalError)?;
let db = session.project_db_mut(&path);
match &path {
AnySystemPath::System(system_path) => {
if let Some(file) = db.files().try_system(db, system_path) {
db.project().close_file(db, file);
} else {
// This can only fail when the path is a directory or it doesn't exists but the
// file should exists for this handler in this branch. This is because every
// close call is preceded by an open call, which ensures that the file is
// interned in the lookup table (`Files`).
tracing::warn!("Salsa file does not exists for {}", system_path);
}
// For non-virtual files, we clear diagnostics if:
//
// 1. The file does not belong to any workspace e.g., opening a random file from
// outside the workspace because closing it acts like the file doesn't exists
// 2. The diagnostic mode is set to open-files only
if session.workspaces().for_path(system_path).is_none()
|| session
.global_settings()
.diagnostic_mode()
.is_open_files_only()
{
clear_diagnostics(session, &url, client);
}
}
AnySystemPath::SystemVirtual(virtual_path) => {
if let Some(virtual_file) = db.files().try_virtual_file(virtual_path) {
db.project().close_file(db, virtual_file.file());
virtual_file.close(db);
} else {
tracing::warn!("Salsa virtual file does not exists for {}", virtual_path);
}
// Always clear diagnostics for virtual files, as they don't really exist on disk
// which means closing them is like deleting the file.
clear_diagnostics(session, &url, client);
}
if should_clear_diagnostics {
clear_diagnostics_if_needed(&document, session, client);
}
Ok(())

15
crates/ty_server/src/server/api/notifications/did_close_notebook.rs
View file

@ -6,8 +6,6 @@ use crate::server::api::LSPResult;
use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
use crate::session::Session;
use crate::session::client::Client;
use crate::system::AnySystemPath;
use ty_project::watch::ChangeEvent;
pub(crate) struct DidCloseNotebookHandler;
@ -30,19 +28,12 @@ impl SyncNotificationHandler for DidCloseNotebookHandler {
.document_handle(&uri)
.with_failure_code(lsp_server::ErrorCode::InternalError)?;
let path = document.to_file_path().into_owned();
document
// We don't need to call publish any diagnostics because we clear
// the diagnostics when closing the corresponding cell documents.
let _ = document
.close(session)
.with_failure_code(lsp_server::ErrorCode::InternalError)?;
if let AnySystemPath::SystemVirtual(virtual_path) = &path {
session.apply_changes(
&path,
vec![ChangeEvent::DeletedVirtual(virtual_path.clone())],
);
}
Ok(())
}
}

46
crates/ty_server/src/server/api/notifications/did_open.rs
View file

@ -1,17 +1,12 @@
use lsp_types::notification::DidOpenTextDocument;
use lsp_types::{DidOpenTextDocumentParams, TextDocumentItem};
use ruff_db::Db as _;
use ruff_db::files::system_path_to_file;
use ty_project::Db as _;
use ty_project::watch::{ChangeEvent, CreatedKind};
use crate::TextDocument;
use crate::server::Result;
use crate::server::api::diagnostics::publish_diagnostics;
use crate::server::api::diagnostics::publish_diagnostics_if_needed;
use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
use crate::session::Session;
use crate::session::client::Client;
use crate::system::AnySystemPath;
pub(crate) struct DidOpenTextDocumentHandler;
@ -39,44 +34,7 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
TextDocument::new(uri, text, version).with_language_id(&language_id),
);
let path = document.to_file_path();
// This is a "maybe" because the `File` might've not been interned yet i.e., the
// `try_system` call will return `None` which doesn't mean that the file is new, it's just
// that the server didn't need the file yet.
let is_maybe_new_system_file = path.as_system().is_some_and(|system_path| {
let db = session.project_db(&path);
db.files()
.try_system(db, system_path)
.is_none_or(|file| !file.exists(db))
});
match &*path {
AnySystemPath::System(system_path) => {
let event = if is_maybe_new_system_file {
ChangeEvent::Created {
path: system_path.clone(),
kind: CreatedKind::File,
}
} else {
ChangeEvent::Opened(system_path.clone())
};
session.apply_changes(&path, vec![event]);
let db = session.project_db_mut(&path);
match system_path_to_file(db, system_path) {
Ok(file) => db.project().open_file(db, file),
Err(err) => tracing::warn!("Failed to open file {system_path}: {err}"),
}
}
AnySystemPath::SystemVirtual(virtual_path) => {
let db = session.project_db_mut(&path);
let virtual_file = db.files().virtual_file(db, virtual_path);
db.project().open_file(db, virtual_file.file());
}
}
publish_diagnostics(session, document.url(), client);
publish_diagnostics_if_needed(&document, session, client);
Ok(())
}

35
crates/ty_server/src/server/api/notifications/did_open_notebook.rs
View file

@ -2,16 +2,14 @@ use lsp_server::ErrorCode;
use lsp_types::DidOpenNotebookDocumentParams;
use lsp_types::notification::DidOpenNotebookDocument;
use ruff_db::Db;
use ty_project::watch::ChangeEvent;
use crate::TextDocument;
use crate::document::NotebookDocument;
use crate::server::Result;
use crate::server::api::LSPResult;
use crate::server::api::diagnostics::publish_diagnostics;
use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
use crate::session::Session;
use crate::session::client::Client;
use crate::system::AnySystemPath;
pub(crate) struct DidOpenNotebookHandler;
@ -22,7 +20,7 @@ impl NotificationHandler for DidOpenNotebookHandler {
impl SyncNotificationHandler for DidOpenNotebookHandler {
fn run(
session: &mut Session,
_client: &Client,
client: &Client,
params: DidOpenNotebookDocumentParams,
) -> Result<()> {
let lsp_types::NotebookDocument {
@ -33,29 +31,22 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
..
} = params.notebook_document;
let notebook = NotebookDocument::new(
notebook_uri,
version,
cells,
metadata.unwrap_or_default(),
params.cell_text_documents,
)
let notebook =
NotebookDocument::new(notebook_uri, version, cells, metadata.unwrap_or_default())
.with_failure_code(ErrorCode::InternalError)?;
let document = session.open_notebook_document(notebook);
let path = document.to_file_path();
let notebook_path = document.notebook_or_file_path();
match &*path {
AnySystemPath::System(system_path) => {
session.apply_changes(&path, vec![ChangeEvent::Opened(system_path.clone())]);
}
AnySystemPath::SystemVirtual(virtual_path) => {
let db = session.project_db_mut(&path);
db.files().virtual_file(db, virtual_path);
}
for cell in params.cell_text_documents {
let cell_document = TextDocument::new(cell.uri, cell.text, cell.version)
.with_language_id(&cell.language_id)
.with_notebook(notebook_path.clone());
session.open_text_document(cell_document);
}
// TODO(dhruvmanila): Publish diagnostics if the client doesn't support pull diagnostics
// Always publish diagnostics because notebooks only support publish diagnostics.
publish_diagnostics(&document, session, client);
Ok(())
}

3
crates/ty_server/src/server/api/requests/completion.rs
View file

@ -44,7 +44,7 @@ impl BackgroundDocumentRequestHandler for CompletionRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};
@ -56,7 +56,6 @@ impl BackgroundDocumentRequestHandler for CompletionRequestHandler {
) else {
return Ok(None);
};
let settings = CompletionSettings {
auto_import: snapshot.global_settings().is_auto_import_enabled(),
};

2
crates/ty_server/src/server/api/requests/diagnostic.rs
View file

@ -33,7 +33,7 @@ impl BackgroundDocumentRequestHandler for DocumentDiagnosticRequestHandler {
_client: &Client,
params: DocumentDiagnosticParams,
) -> Result<DocumentDiagnosticReportResult> {
let diagnostics = compute_diagnostics(db, snapshot);
let diagnostics = compute_diagnostics(db, snapshot.document(), snapshot.encoding());
let Some(diagnostics) = diagnostics else {
return Ok(DocumentDiagnosticReportResult::Report(

2
crates/ty_server/src/server/api/requests/doc_highlights.rs
View file

@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for DocumentHighlightRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/document_symbols.rs
View file

@ -39,7 +39,7 @@ impl BackgroundDocumentRequestHandler for DocumentSymbolRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/goto_declaration.rs
View file

@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for GotoDeclarationRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/goto_definition.rs
View file

@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for GotoDefinitionRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/goto_references.rs
View file

@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for ReferencesRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/goto_type_definition.rs
View file

@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for GotoTypeDefinitionRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/hover.rs
View file

@ -35,7 +35,7 @@ impl BackgroundDocumentRequestHandler for HoverRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

11
crates/ty_server/src/server/api/requests/inlay_hints.rs
View file

@ -1,15 +1,16 @@
use std::borrow::Cow;
use lsp_types::request::InlayHintRequest;
use lsp_types::{InlayHintParams, Url};
use ty_ide::{InlayHintKind, InlayHintLabel, inlay_hints};
use ty_project::ProjectDatabase;
use crate::document::{RangeExt, TextSizeExt};
use crate::server::api::traits::{
BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
};
use crate::session::DocumentSnapshot;
use crate::session::client::Client;
use lsp_types::request::InlayHintRequest;
use lsp_types::{InlayHintParams, Url};
use ty_ide::{InlayHintKind, InlayHintLabel, inlay_hints};
use ty_project::ProjectDatabase;
pub(crate) struct InlayHintRequestHandler;
@ -35,7 +36,7 @@ impl BackgroundDocumentRequestHandler for InlayHintRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/prepare_rename.rs
View file

@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for PrepareRenameRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/rename.rs
View file

@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for RenameRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/selection_range.rs
View file

@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for SelectionRangeRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

26
crates/ty_server/src/server/api/requests/semantic_tokens.rs
View file

@ -1,13 +1,16 @@
use std::borrow::Cow;
use lsp_types::{SemanticTokens, SemanticTokensParams, SemanticTokensResult, Url};
use ruff_db::source::source_text;
use ty_project::ProjectDatabase;
use crate::db::Db;
use crate::server::api::semantic_tokens::generate_semantic_tokens;
use crate::server::api::traits::{
BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
};
use crate::session::DocumentSnapshot;
use crate::session::client::Client;
use lsp_types::{SemanticTokens, SemanticTokensParams, SemanticTokensResult, Url};
use ty_project::ProjectDatabase;
pub(crate) struct SemanticTokensRequestHandler;
@ -33,14 +36,29 @@ impl BackgroundDocumentRequestHandler for SemanticTokensRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};
// If this document is a notebook cell, limit the highlighting range
// to the lines of this cell (instead of highlighting the entire notebook).
// Not only avoids this unnecessary work, this is also required
// because all ranges in the response must be within this **this document**.
let mut cell_range = None;
if snapshot.document().is_cell()
&& let Some(notebook_document) = db.notebook_document(file)
&& let Some(notebook) = source_text(db, file).as_notebook()
{
let cell_index = notebook_document.cell_index_by_uri(snapshot.url());
cell_range = cell_index.and_then(|index| notebook.cell_range(index));
}
let lsp_tokens = generate_semantic_tokens(
db,
file,
None,
cell_range,
snapshot.encoding(),
snapshot
.resolved_client_capabilities()

7
crates/ty_server/src/server/api/requests/semantic_tokens_range.rs
View file

@ -1,5 +1,8 @@
use std::borrow::Cow;
use lsp_types::{SemanticTokens, SemanticTokensRangeParams, SemanticTokensRangeResult, Url};
use ty_project::ProjectDatabase;
use crate::document::RangeExt;
use crate::server::api::semantic_tokens::generate_semantic_tokens;
use crate::server::api::traits::{
@ -7,8 +10,6 @@ use crate::server::api::traits::{
};
use crate::session::DocumentSnapshot;
use crate::session::client::Client;
use lsp_types::{SemanticTokens, SemanticTokensRangeParams, SemanticTokensRangeResult, Url};
use ty_project::ProjectDatabase;
pub(crate) struct SemanticTokensRangeRequestHandler;
@ -34,7 +35,7 @@ impl BackgroundDocumentRequestHandler for SemanticTokensRangeRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/signature_help.rs
View file

@ -38,7 +38,7 @@ impl BackgroundDocumentRequestHandler for SignatureHelpRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.to_file(db) else {
let Some(file) = snapshot.to_notebook_or_file(db) else {
return Ok(None);
};

48
crates/ty_server/src/server/api/requests/workspace_diagnostic.rs
View file

@ -1,3 +1,23 @@
use std::collections::BTreeMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};
use lsp_server::RequestId;
use lsp_types::request::WorkspaceDiagnosticRequest;
use lsp_types::{
FullDocumentDiagnosticReport, PreviousResultId, ProgressToken,
UnchangedDocumentDiagnosticReport, Url, WorkspaceDiagnosticParams, WorkspaceDiagnosticReport,
WorkspaceDiagnosticReportPartialResult, WorkspaceDiagnosticReportResult,
WorkspaceDocumentDiagnosticReport, WorkspaceFullDocumentDiagnosticReport,
WorkspaceUnchangedDocumentDiagnosticReport, notification::Notification,
};
use ruff_db::diagnostic::Diagnostic;
use ruff_db::files::File;
use ruff_db::source::source_text;
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use ty_project::{ProgressReporter, ProjectDatabase};
use crate::PositionEncoding;
use crate::document::DocumentKey;
use crate::server::api::diagnostics::{Diagnostics, to_lsp_diagnostic};
@ -10,23 +30,6 @@ use crate::session::client::Client;
use crate::session::index::Index;
use crate::session::{SessionSnapshot, SuspendedWorkspaceDiagnosticRequest};
use crate::system::file_to_url;
use lsp_server::RequestId;
use lsp_types::request::WorkspaceDiagnosticRequest;
use lsp_types::{
FullDocumentDiagnosticReport, PreviousResultId, ProgressToken,
UnchangedDocumentDiagnosticReport, Url, WorkspaceDiagnosticParams, WorkspaceDiagnosticReport,
WorkspaceDiagnosticReportPartialResult, WorkspaceDiagnosticReportResult,
WorkspaceDocumentDiagnosticReport, WorkspaceFullDocumentDiagnosticReport,
WorkspaceUnchangedDocumentDiagnosticReport, notification::Notification,
};
use ruff_db::diagnostic::Diagnostic;
use ruff_db::files::File;
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};
use ty_project::{ProgressReporter, ProjectDatabase};
/// Handler for [Workspace diagnostics](workspace-diagnostics)
///
@ -368,6 +371,15 @@ impl<'a> ResponseWriter<'a> {
tracing::debug!("Failed to convert file path to URL at {}", file.path(db));
return;
};
if source_text(db, file).is_notebook() {
// Notebooks only support publish diagnostics.
// and we can't convert text ranges to notebook ranges unless
// the document is open in the editor, in which case
// we publish the diagnostics already.
return;
}
let key = DocumentKey::from_url(&url);
let version = self
.index
@ -394,7 +406,7 @@ impl<'a> ResponseWriter<'a> {
new_id => {
let lsp_diagnostics = diagnostics
.iter()
.map(|diagnostic| to_lsp_diagnostic(db, diagnostic, self.position_encoding))
.map(|diagnostic| to_lsp_diagnostic(db, diagnostic, self.position_encoding).1)
.collect::<Vec<_>>();
WorkspaceDocumentDiagnosticReport::Full(WorkspaceFullDocumentDiagnosticReport {

5
crates/ty_server/src/server/api/symbols.rs
View file

@ -26,6 +26,9 @@ pub(crate) fn convert_symbol_kind(kind: ty_ide::SymbolKind) -> SymbolKind {
}
/// Convert a `ty_ide` `SymbolInfo` to LSP `SymbolInformation`
///
/// Returns `None` if the symbol's range cannot be converted to a location
/// (e.g., if the file cannot be converted to a URL).
pub(crate) fn convert_to_lsp_symbol_information(
db: &dyn Db,
file: ruff_db::files::File,
@ -33,12 +36,10 @@ pub(crate) fn convert_to_lsp_symbol_information(
encoding: PositionEncoding,
) -> Option<SymbolInformation> {
let symbol_kind = convert_symbol_kind(symbol.kind);
let location = symbol
.full_range
.to_lsp_range(db, file, encoding)?
.to_location()?;
Some(SymbolInformation {
name: symbol.name.into_owned(),
kind: symbol_kind,

300
crates/ty_server/src/session.rs
View file

@ -1,7 +1,11 @@
//! Data model, state management, and configuration resolution.
use std::collections::{BTreeMap, HashSet, VecDeque};
use std::ops::{Deref, DerefMut};
use std::panic::RefUnwindSafe;
use std::sync::Arc;
use anyhow::{Context, anyhow};
use index::DocumentError;
use lsp_server::{Message, RequestId};
use lsp_types::notification::{DidChangeWatchedFiles, Exit, Notification};
use lsp_types::request::{
@ -13,20 +17,17 @@ use lsp_types::{
DidChangeWatchedFilesRegistrationOptions, FileSystemWatcher, Registration, RegistrationParams,
TextDocumentContentChangeEvent, Unregistration, UnregistrationParams, Url,
};
use options::GlobalOptions;
use ruff_db::Db;
use ruff_db::files::{File, system_path_to_file};
use ruff_db::system::{System, SystemPath, SystemPathBuf};
use std::borrow::Cow;
use std::collections::{BTreeMap, HashSet, VecDeque};
use std::ops::{Deref, DerefMut};
use std::panic::RefUnwindSafe;
use std::sync::Arc;
use ty_combine::Combine;
use ty_project::metadata::Options;
use ty_project::watch::ChangeEvent;
use ty_project::watch::{ChangeEvent, CreatedKind};
use ty_project::{ChangeResult, CheckMode, Db as _, ProjectDatabase, ProjectMetadata};
use index::DocumentError;
use options::GlobalOptions;
pub(crate) use self::options::InitializationOptions;
pub use self::options::{ClientOptions, DiagnosticMode};
pub(crate) use self::settings::{GlobalSettings, WorkspaceSettings};
@ -36,6 +37,7 @@ use crate::capabilities::{
use crate::document::{DocumentKey, DocumentVersion, NotebookDocument};
use crate::server::{Action, publish_settings_diagnostics};
use crate::session::client::Client;
use crate::session::index::Document;
use crate::session::request_queue::RequestQueue;
use crate::system::{AnySystemPath, LSPSystem};
use crate::{PositionEncoding, TextDocument};
@ -816,25 +818,16 @@ impl Session {
let index = self.index();
let document_handle = index.document_handle(url)?;
let notebook = if let Some(notebook_path) = &document_handle.notebook_path {
index
.notebook_arc(&DocumentKey::from(notebook_path.clone()))
.ok()
} else {
None
};
Ok(DocumentSnapshot {
resolved_client_capabilities: self.resolved_client_capabilities,
global_settings: self.global_settings.clone(),
workspace_settings: document_handle
.to_file_path()
.notebook_or_file_path()
.as_system()
.and_then(|path| self.workspaces.settings_for_path(path))
.unwrap_or_else(|| Arc::new(WorkspaceSettings::default())),
position_encoding: self.position_encoding,
document: document_handle,
notebook,
})
}
@ -860,13 +853,7 @@ impl Session {
pub(super) fn text_document_handles(&self) -> impl Iterator<Item = DocumentHandle> + '_ {
self.index()
.text_documents()
.map(|(key, document)| DocumentHandle {
key: key.clone(),
url: document.url().clone(),
version: document.version(),
// TODO: Set notebook path if text document is part of a notebook
notebook_path: None,
})
.map(|(_, document)| DocumentHandle::from_text_document(document))
}
/// Returns a handle to the document specified by its URL.
@ -887,7 +874,7 @@ impl Session {
/// Returns a handle to the opened document.
pub(crate) fn open_notebook_document(&mut self, document: NotebookDocument) -> DocumentHandle {
let handle = self.index_mut().open_notebook_document(document);
self.bump_revision();
self.open_document_in_db(&handle);
handle
}
@ -897,9 +884,49 @@ impl Session {
/// Returns a handle to the opened document.
pub(crate) fn open_text_document(&mut self, document: TextDocument) -> DocumentHandle {
let handle = self.index_mut().open_text_document(document);
self.open_document_in_db(&handle);
handle
}
fn open_document_in_db(&mut self, document: &DocumentHandle) {
let path = document.notebook_or_file_path();
// This is a "maybe" because the `File` might've not been interned yet i.e., the
// `try_system` call will return `None` which doesn't mean that the file is new, it's just
// that the server didn't need the file yet.
let is_maybe_new_system_file = path.as_system().is_some_and(|system_path| {
let db = self.project_db(path);
db.files()
.try_system(db, system_path)
.is_none_or(|file| !file.exists(db))
});
match path {
AnySystemPath::System(system_path) => {
let event = if is_maybe_new_system_file {
ChangeEvent::Created {
path: system_path.clone(),
kind: CreatedKind::File,
}
} else {
ChangeEvent::Opened(system_path.clone())
};
self.apply_changes(path, vec![event]);
let db = self.project_db_mut(path);
match system_path_to_file(db, system_path) {
Ok(file) => db.project().open_file(db, file),
Err(err) => tracing::warn!("Failed to open file {system_path}: {err}"),
}
}
AnySystemPath::SystemVirtual(virtual_path) => {
let db = self.project_db_mut(path);
let virtual_file = db.files().virtual_file(db, virtual_path);
db.project().open_file(db, virtual_file.file());
}
}
self.bump_revision();
handle
}
/// Returns a reference to the index.
@ -999,7 +1026,6 @@ pub(crate) struct DocumentSnapshot {
workspace_settings: Arc<WorkspaceSettings>,
position_encoding: PositionEncoding,
document: DocumentHandle,
notebook: Option<Arc<NotebookDocument>>,
}
impl DocumentSnapshot {
@ -1028,17 +1054,12 @@ impl DocumentSnapshot {
&self.document
}
/// Returns the URL of the document.
pub(crate) fn url(&self) -> &lsp_types::Url {
self.document.url()
}
pub(crate) fn notebook(&self) -> Option<&NotebookDocument> {
self.notebook.as_deref()
}
pub(crate) fn to_file(&self, db: &dyn Db) -> Option<File> {
let file = self.document.to_file(db);
pub(crate) fn to_notebook_or_file(&self, db: &dyn Db) -> Option<File> {
let file = self.document.notebook_or_file(db);
if file.is_none() {
tracing::debug!(
"Failed to resolve file: file not found for `{}`",
@ -1048,8 +1069,8 @@ impl DocumentSnapshot {
file
}
pub(crate) fn to_file_path(&self) -> Cow<'_, AnySystemPath> {
self.document.to_file_path()
pub(crate) fn notebook_or_file_path(&self) -> &AnySystemPath {
self.document.notebook_or_file_path()
}
}
@ -1330,34 +1351,99 @@ impl SuspendedWorkspaceDiagnosticRequest {
///
/// It also exposes methods to get the file-path of the corresponding ty-file.
#[derive(Clone, Debug)]
pub(crate) struct DocumentHandle {
/// The key that uniquely identifies this document in the index.
key: DocumentKey,
pub(crate) enum DocumentHandle {
Text {
url: lsp_types::Url,
/// The path to the enclosing notebook file if this document is a notebook or a notebook cell.
notebook_path: Option<AnySystemPath>,
path: AnySystemPath,
version: DocumentVersion,
},
Notebook {
url: lsp_types::Url,
path: AnySystemPath,
version: DocumentVersion,
},
Cell {
url: lsp_types::Url,
version: DocumentVersion,
notebook_path: AnySystemPath,
},
}
impl DocumentHandle {
fn from_text_document(document: &TextDocument) -> Self {
match document.notebook() {
None => Self::Text {
version: document.version(),
url: document.url().clone(),
path: DocumentKey::from_url(document.url()).into_file_path(),
},
Some(notebook) => Self::Cell {
notebook_path: notebook.clone(),
version: document.version(),
url: document.url().clone(),
},
}
}
fn from_notebook_document(document: &NotebookDocument) -> Self {
Self::Notebook {
path: DocumentKey::from_url(document.url()).into_file_path(),
url: document.url().clone(),
version: document.version(),
}
}
fn from_document(document: &Document) -> Self {
match document {
Document::Text(text) => Self::from_text_document(text),
Document::Notebook(notebook) => Self::from_notebook_document(notebook),
}
}
fn key(&self) -> DocumentKey {
DocumentKey::from_url(self.url())
}
pub(crate) const fn version(&self) -> DocumentVersion {
self.version
match self {
Self::Text { version, .. }
| Self::Notebook { version, .. }
| Self::Cell { version, .. } => *version,
}
}
/// The URL as used by the client to reference this document.
pub(crate) fn url(&self) -> &lsp_types::Url {
&self.url
match self {
Self::Text { url, .. } | Self::Notebook { url, .. } | Self::Cell { url, .. } => url,
}
}
/// The path to the enclosing file for this document.
///
/// This is the path corresponding to the URL, except for notebook cells where the
/// path corresponds to the notebook file.
pub(crate) fn to_file_path(&self) -> Cow<'_, AnySystemPath> {
if let Some(path) = self.notebook_path.as_ref() {
Cow::Borrowed(path)
} else {
Cow::Owned(self.key.to_file_path())
pub(crate) fn notebook_or_file_path(&self) -> &AnySystemPath {
match self {
Self::Text { path, .. } | Self::Notebook { path, .. } => path,
Self::Cell { notebook_path, .. } => notebook_path,
}
}
#[expect(unused)]
pub(crate) fn file_path(&self) -> Option<&AnySystemPath> {
match self {
Self::Text { path, .. } | Self::Notebook { path, .. } => Some(path),
Self::Cell { .. } => None,
}
}
#[expect(unused)]
pub(crate) fn notebook_path(&self) -> Option<&AnySystemPath> {
match self {
DocumentHandle::Notebook { path, .. } => Some(path),
DocumentHandle::Cell { notebook_path, .. } => Some(notebook_path),
DocumentHandle::Text { .. } => None,
}
}
@ -1366,8 +1452,8 @@ impl DocumentHandle {
/// It returns [`None`] for the following cases:
/// - For virtual file, if it's not yet opened
/// - For regular file, if it does not exists or is a directory
pub(crate) fn to_file(&self, db: &dyn Db) -> Option<File> {
match &*self.to_file_path() {
pub(crate) fn notebook_or_file(&self, db: &dyn Db) -> Option<File> {
match &self.notebook_or_file_path() {
AnySystemPath::System(path) => system_path_to_file(db, path).ok(),
AnySystemPath::SystemVirtual(virtual_path) => db
.files()
@ -1376,6 +1462,14 @@ impl DocumentHandle {
}
}
pub(crate) fn is_cell(&self) -> bool {
matches!(self, Self::Cell { .. })
}
pub(crate) fn is_cell_or_notebook(&self) -> bool {
matches!(self, Self::Cell { .. } | Self::Notebook { .. })
}
pub(crate) fn update_text_document(
&self,
session: &mut Session,
@ -1383,9 +1477,10 @@ impl DocumentHandle {
new_version: DocumentVersion,
) -> crate::Result<()> {
let position_encoding = session.position_encoding();
{
let mut index = session.index_mut();
let document_mut = index.document_mut(&self.key)?;
let document_mut = index.document_mut(&self.key())?;
let Some(document) = document_mut.as_text_mut() else {
anyhow::bail!("Text document path does not point to a text document");
@ -1393,19 +1488,110 @@ impl DocumentHandle {
if content_changes.is_empty() {
document.update_version(new_version);
return Ok(());
} else {
document.apply_changes(content_changes, new_version, position_encoding);
}
}
document.apply_changes(content_changes, new_version, position_encoding);
self.update_in_db(session);
Ok(())
}
pub(crate) fn update_notebook_document(
&self,
session: &mut Session,
cells: Option<lsp_types::NotebookDocumentCellChange>,
metadata: Option<lsp_types::LSPObject>,
new_version: DocumentVersion,
) -> crate::Result<()> {
let position_encoding = session.position_encoding();
{
let mut index = session.index_mut();
index.update_notebook_document(
&self.key(),
cells,
metadata,
new_version,
position_encoding,
)?;
}
self.update_in_db(session);
Ok(())
}
fn update_in_db(&self, session: &mut Session) {
let path = self.notebook_or_file_path();
let changes = match path {
AnySystemPath::System(system_path) => {
vec![ChangeEvent::file_content_changed(system_path.clone())]
}
AnySystemPath::SystemVirtual(virtual_path) => {
vec![ChangeEvent::ChangedVirtual(virtual_path.clone())]
}
};
session.apply_changes(path, changes);
}
/// De-registers a document, specified by its key.
/// Calling this multiple times for the same document is a logic error.
pub(crate) fn close(self, session: &mut Session) -> crate::Result<()> {
session.index_mut().close_document(&self.key)?;
///
/// Returns `true` if the client needs to clear the diagnostics for this document.
pub(crate) fn close(&self, session: &mut Session) -> crate::Result<bool> {
let is_cell = self.is_cell();
let path = self.notebook_or_file_path();
session.index_mut().close_document(&self.key())?;
// Close the text or notebook file in the database but skip this
// step for cells because closing a cell doesn't close its notebook.
let requires_clear_diagnostics = if is_cell {
true
} else {
let db = session.project_db_mut(path);
match path {
AnySystemPath::System(system_path) => {
if let Some(file) = db.files().try_system(db, system_path) {
db.project().close_file(db, file);
} else {
// This can only fail when the path is a directory or it doesn't exists but the
// file should exists for this handler in this branch. This is because every
// close call is preceded by an open call, which ensures that the file is
// interned in the lookup table (`Files`).
tracing::warn!("Salsa file does not exists for {}", system_path);
}
// For non-virtual files, we clear diagnostics if:
//
// 1. The file does not belong to any workspace e.g., opening a random file from
// outside the workspace because closing it acts like the file doesn't exists
// 2. The diagnostic mode is set to open-files only
session.workspaces().for_path(system_path).is_none()
|| session
.global_settings()
.diagnostic_mode()
.is_open_files_only()
}
AnySystemPath::SystemVirtual(virtual_path) => {
if let Some(virtual_file) = db.files().try_virtual_file(virtual_path) {
db.project().close_file(db, virtual_file.file());
virtual_file.close(db);
} else {
tracing::warn!("Salsa virtual file does not exists for {}", virtual_path);
}
// Always clear diagnostics for virtual files, as they don't really exist on disk
// which means closing them is like deleting the file.
true
}
}
};
session.bump_revision();
Ok(())
Ok(requires_clear_diagnostics)
}
}

185
crates/ty_server/src/session/index.rs
View file

@ -1,3 +1,4 @@
use rustc_hash::FxHashMap;
use std::sync::Arc;
use crate::document::DocumentKey;
@ -5,27 +6,19 @@ use crate::session::DocumentHandle;
use crate::{
PositionEncoding, TextDocument,
document::{DocumentVersion, NotebookDocument},
system::AnySystemPath,
};
use ruff_db::system::SystemVirtualPath;
use rustc_hash::FxHashMap;
/// Stores and tracks all open documents in a session, along with their associated settings.
#[derive(Debug)]
pub(crate) struct Index {
/// Maps all document file paths to the associated document controller
documents: FxHashMap<DocumentKey, Document>,
/// Maps opaque cell URLs to a notebook path (document)
notebook_cells: FxHashMap<String, AnySystemPath>,
}
impl Index {
pub(super) fn new() -> Self {
Self {
documents: FxHashMap::default(),
notebook_cells: FxHashMap::default(),
}
}
@ -47,23 +40,7 @@ impl Index {
return Err(DocumentError::NotFound(key));
};
if let Some(path) = key.as_opaque() {
if let Some(notebook_path) = self.notebook_cells.get(path) {
return Ok(DocumentHandle {
key: key.clone(),
notebook_path: Some(notebook_path.clone()),
url: url.clone(),
version: document.version(),
});
}
}
Ok(DocumentHandle {
key: key.clone(),
notebook_path: None,
url: url.clone(),
version: document.version(),
})
Ok(DocumentHandle::from_document(document))
}
#[expect(dead_code)]
@ -74,7 +51,6 @@ impl Index {
.map(|(key, _)| key)
}
#[expect(dead_code)]
pub(super) fn update_notebook_document(
&mut self,
notebook_key: &DocumentKey,
@ -83,26 +59,100 @@ impl Index {
new_version: DocumentVersion,
encoding: PositionEncoding,
) -> crate::Result<()> {
// update notebook cell index
if let Some(lsp_types::NotebookDocumentCellChangeStructure {
did_open: Some(did_open),
..
}) = cells.as_ref().and_then(|cells| cells.structure.as_ref())
{
for opened_cell in did_open {
let cell_path = SystemVirtualPath::new(opened_cell.uri.as_str());
self.notebook_cells
.insert(cell_path.to_string(), notebook_key.to_file_path());
}
// deleted notebook cells are closed via textDocument/didClose - we don't close them here.
}
let document = self.document_mut(notebook_key)?;
let Some(notebook) = document.as_notebook_mut() else {
anyhow::bail!("Notebook document path does not point to a notebook document");
};
notebook.update(cells, metadata, new_version, encoding)?;
let (structure, data, text_content) = cells
.map(|cells| {
let lsp_types::NotebookDocumentCellChange {
structure,
data,
text_content,
} = cells;
(structure, data, text_content)
})
.unwrap_or_default();
let (array, did_open, did_close) = structure
.map(|structure| {
let lsp_types::NotebookDocumentCellChangeStructure {
array,
did_open,
did_close,
} = structure;
(array, did_open, did_close)
})
.unwrap_or_else(|| {
(
lsp_types::NotebookCellArrayChange {
start: 0,
delete_count: 0,
cells: None,
},
None,
None,
)
});
tracing::info!(
"version: {}, new_version: {}",
notebook.version(),
new_version
);
notebook.update(array, data.unwrap_or_default(), metadata, new_version)?;
let notebook_path = notebook_key.to_file_path();
for opened_cell in did_open.into_iter().flatten() {
self.documents.insert(
DocumentKey::from_url(&opened_cell.uri),
Document::Text(
TextDocument::new(opened_cell.uri, opened_cell.text, opened_cell.version)
.with_language_id(&opened_cell.language_id)
.with_notebook(notebook_path.clone())
.into(),
),
);
}
for updated_cell in text_content.into_iter().flatten() {
let Ok(document_mut) =
self.document_mut(&DocumentKey::from_url(&updated_cell.document.uri))
else {
tracing::warn!(
"Could not find document for cell {}",
updated_cell.document.uri
);
continue;
};
let Some(document) = document_mut.as_text_mut() else {
continue;
};
if updated_cell.changes.is_empty() {
document.update_version(updated_cell.document.version);
} else {
document.apply_changes(
updated_cell.changes,
updated_cell.document.version,
encoding,
);
}
}
// VS Code sends a separate `didClose` request for every cell
// and they're removed from the metadata (notebook document)
// because they get deleted as part of `change.cells.structure.array`
let _ = did_close;
let notebook = self.document(notebook_key).unwrap().as_notebook().unwrap();
let ruff_notebook = notebook.to_ruff_notebook(self);
tracing::debug!("Updated notebook: {:?}", ruff_notebook.source_code());
Ok(())
}
@ -117,31 +167,10 @@ impl Index {
Ok(document)
}
pub(crate) fn notebook_arc(
&self,
key: &DocumentKey,
) -> Result<Arc<NotebookDocument>, DocumentError> {
let Some(document) = self.documents.get(key) else {
return Err(DocumentError::NotFound(key.clone()));
};
if let Document::Notebook(notebook) = document {
Ok(notebook.clone())
} else {
Err(DocumentError::NotFound(key.clone()))
}
}
pub(super) fn open_text_document(&mut self, document: TextDocument) -> DocumentHandle {
let key = DocumentKey::from_url(document.url());
// TODO: Fix file path for notebook cells
let handle = DocumentHandle {
key: key.clone(),
notebook_path: None,
url: document.url().clone(),
version: document.version(),
};
let handle = DocumentHandle::from_text_document(&document);
self.documents.insert(key, Document::new_text(document));
@ -149,38 +178,18 @@ impl Index {
}
pub(super) fn open_notebook_document(&mut self, document: NotebookDocument) -> DocumentHandle {
let handle = DocumentHandle::from_notebook_document(&document);
let notebook_key = DocumentKey::from_url(document.url());
let url = document.url().clone();
let version = document.version();
for cell_url in document.cell_urls() {
self.notebook_cells
.insert(cell_url.to_string(), notebook_key.to_file_path());
}
self.documents
.insert(notebook_key.clone(), Document::new_notebook(document));
.insert(notebook_key, Document::new_notebook(document));
DocumentHandle {
notebook_path: Some(notebook_key.to_file_path()),
key: notebook_key,
url,
version,
}
}
pub(super) fn close_document(&mut self, key: &DocumentKey) -> crate::Result<()> {
// Notebook cells URIs are removed from the index here, instead of during
// `update_notebook_document`. This is because a notebook cell, as a text document,
// is requested to be `closed` by VS Code after the notebook gets updated.
// This is not documented in the LSP specification explicitly, and this assumption
// may need revisiting in the future as we support more editors with notebook support.
if let DocumentKey::Opaque(uri) = key {
self.notebook_cells.remove(uri);
handle
}
pub(super) fn close_document(&mut self, key: &DocumentKey) -> Result<(), DocumentError> {
let Some(_) = self.documents.remove(key) else {
anyhow::bail!("tried to close document that didn't exist at {key}")
return Err(DocumentError::NotFound(key.clone()));
};
Ok(())

65
crates/ty_server/src/system.rs
View file

@ -1,11 +1,12 @@
use std::any::Any;
use std::fmt;
use std::fmt::Display;
use std::hash::{DefaultHasher, Hash, Hasher as _};
use std::panic::RefUnwindSafe;
use std::sync::Arc;
use crate::Db;
use crate::document::DocumentKey;
use crate::document::{DocumentKey, LanguageId};
use crate::session::index::{Document, Index};
use lsp_types::Url;
use ruff_db::file_revision::FileRevision;
@ -16,6 +17,7 @@ use ruff_db::system::{
SystemPath, SystemPathBuf, SystemVirtualPath, SystemVirtualPathBuf, WritableSystem,
};
use ruff_notebook::{Notebook, NotebookError};
use ruff_python_ast::PySourceType;
use ty_ide::cached_vendored_path;
/// Returns a [`Url`] for the given [`File`].
@ -117,6 +119,23 @@ impl LSPSystem {
index.document(&DocumentKey::from(path)).ok()
}
fn source_type_from_document(
document: &Document,
extension: Option<&str>,
) -> Option<PySourceType> {
match document {
Document::Text(text) => match text.language_id()? {
LanguageId::Python => Some(
extension
.and_then(PySourceType::try_from_extension)
.unwrap_or(PySourceType::Python),
),
LanguageId::Other => None,
},
Document::Notebook(_) => Some(PySourceType::Ipynb),
}
}
pub(crate) fn system_path_to_document(&self, path: &SystemPath) -> Option<&Document> {
let any_path = AnySystemPath::System(path.to_path_buf());
self.document(any_path)
@ -137,7 +156,7 @@ impl System for LSPSystem {
if let Some(document) = document {
Ok(Metadata::new(
document_revision(document),
document_revision(document, self.index()),
None,
FileType::File,
))
@ -154,6 +173,16 @@ impl System for LSPSystem {
self.native_system.path_exists_case_sensitive(path, prefix)
}
fn source_type(&self, path: &SystemPath) -> Option<PySourceType> {
let document = self.system_path_to_document(path)?;
Self::source_type_from_document(document, path.extension())
}
fn virtual_path_source_type(&self, path: &SystemVirtualPath) -> Option<PySourceType> {
let document = self.system_virtual_path_to_document(path)?;
Self::source_type_from_document(document, path.extension())
}
fn read_to_string(&self, path: &SystemPath) -> Result<String> {
let document = self.system_path_to_document(path);
@ -168,7 +197,7 @@ impl System for LSPSystem {
match document {
Some(Document::Text(document)) => Notebook::from_source_code(document.contents()),
Some(Document::Notebook(notebook)) => Ok(notebook.make_ruff_notebook()),
Some(Document::Notebook(notebook)) => Ok(notebook.to_ruff_notebook(self.index())),
None => self.native_system.read_to_notebook(path),
}
}
@ -195,7 +224,7 @@ impl System for LSPSystem {
match document {
Document::Text(document) => Notebook::from_source_code(document.contents()),
Document::Notebook(notebook) => Ok(notebook.make_ruff_notebook()),
Document::Notebook(notebook) => Ok(notebook.to_ruff_notebook(self.index())),
}
}
@ -272,9 +301,33 @@ fn virtual_path_not_found(path: impl Display) -> std::io::Error {
}
/// Helper function to get the [`FileRevision`] of the given document.
fn document_revision(document: &Document) -> FileRevision {
fn document_revision(document: &Document, index: &Index) -> FileRevision {
// The file revision is just an opaque number which doesn't have any significant meaning other
// than that the file has changed if the revisions are different.
#[expect(clippy::cast_sign_loss)]
FileRevision::new(document.version() as u128)
match document {
Document::Text(text) => FileRevision::new(text.version() as u128),
Document::Notebook(notebook) => {
// VS Code doesn't always bump the notebook version when the cell content changes.
// Specifically, I noticed that VS Code re-uses the same version when:
// 1. Adding a new cell
// 2. Pasting some code that has an error
//
// The notification updating the cell content on paste re-used the same version as when the cell was added.
// Because of that, hash all cell versions and the notebook versions together.
let mut hasher = DefaultHasher::new();
for cell_url in notebook.cell_urls() {
if let Ok(cell) = index.document(&DocumentKey::from_url(cell_url)) {
cell.version().hash(&mut hasher);
}
}
// Use higher 64 bits for notebook version and lower 64 bits for cell revisions
let notebook_version_high = (notebook.version() as u128) << 64;
let cell_versions_low = u128::from(hasher.finish()) & 0xFFFF_FFFF_FFFF_FFFF;
let combined_revision = notebook_version_high | cell_versions_low;
FileRevision::new(combined_revision)
}
}
}

56
crates/ty_server/tests/e2e/main.rs
View file

@ -30,11 +30,12 @@
mod commands;
mod initialize;
mod inlay_hints;
mod notebook;
mod publish_diagnostics;
mod pull_diagnostics;
use std::collections::hash_map::Entry;
use std::collections::{HashMap, VecDeque};
use std::collections::{BTreeMap, HashMap, VecDeque};
use std::num::NonZeroUsize;
use std::sync::{Arc, OnceLock};
use std::thread::JoinHandle;
@ -433,6 +434,33 @@ impl TestServer {
))
}
/// Collects `N` publish diagnostic notifications into a map, indexed by the document url.
///
/// ## Panics
/// If there are multiple publish diagnostics notifications for the same document.
pub(crate) fn collect_publish_diagnostic_notifications(
&mut self,
count: usize,
) -> Result<BTreeMap<lsp_types::Url, Vec<lsp_types::Diagnostic>>> {
let mut results = BTreeMap::default();
for _ in 0..count {
let notification =
self.await_notification::<lsp_types::notification::PublishDiagnostics>()?;
if let Some(existing) =
results.insert(notification.uri.clone(), notification.diagnostics)
{
panic!(
"Received multiple publish diagnostic notifications for {url}: ({existing:#?})",
url = &notification.uri
);
}
}
Ok(results)
}
/// Wait for a request of the specified type from the server and return the request ID and
/// parameters.
///
@ -774,6 +802,7 @@ impl Drop for TestServer {
match self.await_response::<Shutdown>(&shutdown_id) {
Ok(()) => {
self.send_notification::<Exit>(());
None
}
Err(err) => Some(format!("Failed to get shutdown response: {err:?}")),
@ -782,9 +811,32 @@ impl Drop for TestServer {
None
};
if let Some(_client_connection) = self.client_connection.take() {
// Drop the client connection before joining the server thread to avoid any hangs
// in case the server didn't respond to the shutdown request.
if let Some(client_connection) = self.client_connection.take() {
if !std::thread::panicking() {
// Wait for the client sender to drop (confirmation that it processed the exit notification).
match client_connection
.receiver
.recv_timeout(Duration::from_secs(20))
{
Err(RecvTimeoutError::Disconnected) => {
// Good, the server terminated
}
Err(RecvTimeoutError::Timeout) => {
tracing::warn!(
"The server didn't exit within 20ms after receiving the EXIT notification"
);
}
Ok(message) => {
// Ignore any errors: A duplicate pending message
// won't matter that much because `assert_no_pending_messages` will
// panic anyway.
let _ = self.handle_message(message);
}
}
}
}
if std::thread::panicking() {

361
crates/ty_server/tests/e2e/notebook.rs Normal file
View file

@ -0,0 +1,361 @@
use insta::assert_json_snapshot;
use lsp_types::{NotebookCellKind, Position, Range};
use crate::{TestServer, TestServerBuilder};
#[test]
fn publish_diagnostics_open() -> 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)
"#,
);
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",
"relatedInformation": []
}
],
"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(())
}
fn semantic_tokens_full_for_cell(
server: &mut TestServer,
cell_uri: &lsp_types::Url,
) -> crate::Result<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
}
}

13
crates/ty_server/tests/e2e/snapshots/e2e__initialize__initialization.snap
View file

@ -1,6 +1,5 @@
---
source: crates/ty_server/tests/e2e/initialize.rs
assertion_line: 17
expression: initialization_result
---
{
@ -10,6 +9,18 @@ expression: initialization_result
"openClose": true,
"change": 2
},
"notebookDocumentSync": {
"notebookSelector": [
{
"cells": [
{
"language": "python"
}
]
}
],
"save": false
},
"selectionRangeProvider": true,
"hoverProvider": true,
"completionProvider": {

13
crates/ty_server/tests/e2e/snapshots/e2e__initialize__initialization_with_workspace.snap
View file

@ -1,6 +1,5 @@
---
source: crates/ty_server/tests/e2e/initialize.rs
assertion_line: 32
expression: initialization_result
---
{
@ -10,6 +9,18 @@ expression: initialization_result
"openClose": true,
"change": 2
},
"notebookDocumentSync": {
"notebookSelector": [
{
"cells": [
{
"language": "python"
}
]
}
],
"save": false
},
"selectionRangeProvider": true,
"hoverProvider": true,
"completionProvider": {

263
crates/ty_server/tests/e2e/snapshots/e2e__notebook__diagnostic_end_of_file.snap Normal file
View file

@ -0,0 +1,263 @@
---
source: crates/ty_server/tests/e2e/notebook.rs
expression: diagnostics
---
{
"vscode-notebook-cell://test.ipynb#0": [],
"vscode-notebook-cell://test.ipynb#1": [
{
"range": {
"start": {
"line": 0,
"character": 49
},
"end": {
"line": 0,
"character": 52
}
},
"severity": 1,
"code": "invalid-return-type",
"codeDescription": {
"href": "https://ty.dev/rules#invalid-return-type"
},
"source": "ty",
"message": "Function can implicitly return `None`, which is not assignable to return type `str`",
"relatedInformation": []
}
],
"vscode-notebook-cell://test.ipynb#2": [
{
"range": {
"start": {
"line": 0,
"character": 19
},
"end": {
"line": 0,
"character": 23
}
},
"severity": 1,
"code": "invalid-syntax",
"source": "ty",
"message": "Expected `,`, found name",
"relatedInformation": []
},
{
"range": {
"start": {
"line": 0,
"character": 19
},
"end": {
"line": 0,
"character": 23
}
},
"severity": 1,
"code": "unresolved-reference",
"codeDescription": {
"href": "https://ty.dev/rules#unresolved-reference"
},
"source": "ty",
"message": "Name `word` used when not defined",
"relatedInformation": []
},
{
"range": {
"start": {
"line": 0,
"character": 23
},
"end": {
"line": 0,
"character": 27
}
},
"severity": 1,
"code": "invalid-syntax",
"source": "ty",
"message": "Expected `,`, found string",
"relatedInformation": []
},
{
"range": {
"start": {
"line": 0,
"character": 23
},
"end": {
"line": 0,
"character": 27
}
},
"severity": 1,
"code": "invalid-argument-type",
"codeDescription": {
"href": "https://ty.dev/rules#invalid-argument-type"
},
"source": "ty",
"message": "Argument to function `with_style` is incorrect: Expected `Style`, found `Literal[/", /"]`",
"relatedInformation": [
{
"location": {
"uri": "vscode-notebook-cell://test.ipynb#1",
"range": {
"start": {
"line": 0,
"character": 4
},
"end": {
"line": 0,
"character": 14
}
}
},
"message": "Function defined here"
},
{
"location": {
"uri": "vscode-notebook-cell://test.ipynb#1",
"range": {
"start": {
"line": 0,
"character": 32
},
"end": {
"line": 0,
"character": 44
}
}
},
"message": "Parameter declared here"
}
]
},
{
"range": {
"start": {
"line": 0,
"character": 27
},
"end": {
"line": 0,
"character": 36
}
},
"severity": 1,
"code": "invalid-syntax",
"source": "ty",
"message": "Expected `,`, found name",
"relatedInformation": []
},
{
"range": {
"start": {
"line": 0,
"character": 27
},
"end": {
"line": 0,
"character": 36
}
},
"severity": 1,
"code": "unresolved-reference",
"codeDescription": {
"href": "https://ty.dev/rules#unresolved-reference"
},
"source": "ty",
"message": "Name `underline` used when not defined",
"relatedInformation": []
},
{
"range": {
"start": {
"line": 0,
"character": 27
},
"end": {
"line": 0,
"character": 36
}
},
"severity": 1,
"code": "too-many-positional-arguments",
"codeDescription": {
"href": "https://ty.dev/rules#too-many-positional-arguments"
},
"source": "ty",
"message": "Too many positional arguments to function `with_style`: expected 3, got 6",
"relatedInformation": [
{
"location": {
"uri": "vscode-notebook-cell://test.ipynb#1",
"range": {
"start": {
"line": 0,
"character": 4
},
"end": {
"line": 0,
"character": 52
}
}
},
"message": "Function signature here"
}
]
},
{
"range": {
"start": {
"line": 0,
"character": 36
},
"end": {
"line": 0,
"character": 38
}
},
"severity": 1,
"code": "invalid-syntax",
"source": "ty",
"message": "missing closing quote in string literal",
"relatedInformation": []
},
{
"range": {
"start": {
"line": 2,
"character": 0
},
"end": {
"line": 2,
"character": 7
}
},
"severity": 1,
"code": "invalid-syntax",
"source": "ty",
"message": "Expected `,`, found name",
"relatedInformation": []
},
{
"range": {
"start": {
"line": 3,
"character": 0
},
"end": {
"line": 3,
"character": 0
}
},
"severity": 1,
"code": "invalid-syntax",
"source": "ty",
"message": "unexpected EOF while parsing",
"relatedInformation": []
}
]
}

96
crates/ty_server/tests/e2e/snapshots/e2e__notebook__publish_diagnostics_open.snap Normal file
View file

@ -0,0 +1,96 @@
---
source: crates/ty_server/tests/e2e/notebook.rs
expression: "[cell1_diagnostics, cell2_diagnostics, cell3_diagnostics]"
---
[
{
"uri": "vscode-notebook-cell://test.ipynb#1",
"diagnostics": [
{
"range": {
"start": {
"line": 0,
"character": 49
},
"end": {
"line": 0,
"character": 52
}
},
"severity": 1,
"code": "invalid-return-type",
"codeDescription": {
"href": "https://ty.dev/rules#invalid-return-type"
},
"source": "ty",
"message": "Function can implicitly return `None`, which is not assignable to return type `str`",
"relatedInformation": []
}
],
"version": 0
},
{
"uri": "vscode-notebook-cell://test.ipynb#2",
"diagnostics": [
{
"range": {
"start": {
"line": 0,
"character": 66
},
"end": {
"line": 0,
"character": 78
}
},
"severity": 1,
"code": "invalid-argument-type",
"codeDescription": {
"href": "https://ty.dev/rules#invalid-argument-type"
},
"source": "ty",
"message": "Argument to function `with_style` is incorrect: Expected `Style`, found `Literal[/"underlined/"]`",
"relatedInformation": [
{
"location": {
"uri": "vscode-notebook-cell://test.ipynb#1",
"range": {
"start": {
"line": 0,
"character": 4
},
"end": {
"line": 0,
"character": 14
}
}
},
"message": "Function defined here"
},
{
"location": {
"uri": "vscode-notebook-cell://test.ipynb#1",
"range": {
"start": {
"line": 0,
"character": 32
},
"end": {
"line": 0,
"character": 44
}
}
},
"message": "Parameter declared here"
}
]
}
],
"version": 0
},
{
"uri": "vscode-notebook-cell://test.ipynb#0",
"diagnostics": [],
"version": 0
}
]

263
crates/ty_server/tests/e2e/snapshots/e2e__notebook__semantic_tokens.snap Normal file
View file

@ -0,0 +1,263 @@
---
source: crates/ty_server/tests/e2e/notebook.rs
expression: "[cell1_tokens, cell2_tokens, cell3_tokens]"
---
[
{
"data": [
0,
5,
6,
0,
0,
0,
14,
7,
5,
0,
2,
5,
5,
14,
0,
0,
8,
7,
5,
0,
0,
8,
8,
10,
0,
0,
10,
6,
10,
0,
0,
8,
11,
10,
0
]
},
{
"data": [
0,
4,
10,
7,
1,
0,
11,
4,
2,
1,
0,
6,
3,
1,
0,
0,
5,
4,
2,
1,
0,
6,
5,
2,
1,
0,
7,
5,
14,
0,
0,
10,
3,
1,
0,
1,
7,
5,
2,
0,
0,
9,
8,
10,
0,
1,
15,
4,
2,
0,
0,
5,
7,
8,
0,
0,
8,
4,
2,
0,
0,
8,
1,
10,
0,
0,
2,
4,
2,
0,
0,
5,
1,
10,
0,
1,
9,
5,
2,
0,
0,
9,
6,
10,
0,
1,
15,
4,
2,
0,
0,
5,
7,
8,
0,
0,
8,
4,
2,
0,
0,
8,
2,
10,
0,
0,
3,
4,
2,
0,
0,
5,
2,
10,
0,
2,
4,
8,
5,
1,
0,
11,
4,
2,
0,
0,
5,
4,
8,
0,
0,
5,
4,
2,
0,
1,
4,
6,
5,
1,
0,
9,
4,
2,
0,
0,
7,
4,
10,
0,
1,
4,
6,
5,
0,
0,
10,
3,
10,
0,
0,
6,
8,
5,
0,
1,
4,
6,
5,
0,
0,
10,
3,
10,
0,
0,
6,
3,
7,
0,
0,
4,
4,
2,
0
]
},
{
"data": [
0,
0,
5,
7,
0,
0,
6,
10,
7,
0,
0,
11,
39,
10,
0,
0,
41,
6,
10,
0,
0,
8,
12,
10,
0
]
}
]