Python
/
ruff
mirror of https://github.com/astral-sh/ruff
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[ty] Smaller refactors to server API in prep for notebook support (#21095)

This commit is contained in:
Micha Reiser 2025-10-31 21:00:04 +01:00 committed by GitHub
parent 827d8ae5d4
commit 6337e22f0c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
33 changed files with 570 additions and 546 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
crates/ruff_db/src/system/path.rs
View File

@ -723,10 +723,11 @@ impl ruff_cache::CacheKey for SystemPathBuf {
/// A slice of a virtual path on [`System`](super::System) (akin to [`str`]).
#[repr(transparent)]
#[derive(Eq, PartialEq, Hash, PartialOrd, Ord)]
pub struct SystemVirtualPath(str);
impl SystemVirtualPath {
pub fn new(path: &str) -> &SystemVirtualPath {
pub const fn new(path: &str) -> &SystemVirtualPath {
// SAFETY: SystemVirtualPath is marked as #[repr(transparent)] so the conversion from a
// *const str to a *const SystemVirtualPath is valid.
unsafe { &*(path as *const str as *const SystemVirtualPath) }
@ -767,8 +768,8 @@ pub struct SystemVirtualPathBuf(String);
impl SystemVirtualPathBuf {
#[inline]
pub fn as_path(&self) -> &SystemVirtualPath {
SystemVirtualPath::new(&self.0)
pub const fn as_path(&self) -> &SystemVirtualPath {
SystemVirtualPath::new(self.0.as_str())
}
}
@ -852,6 +853,12 @@ impl ruff_cache::CacheKey for SystemVirtualPathBuf {
}
}
impl Borrow<SystemVirtualPath> for SystemVirtualPathBuf {
fn borrow(&self) -> &SystemVirtualPath {
self.as_path()
}
}
/// Deduplicates identical paths and removes nested paths.
///
/// # Examples

90
crates/ty_server/src/document.rs
View File

@ -11,6 +11,7 @@ use lsp_types::{PositionEncodingKind, Url};
use crate::system::AnySystemPath;
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;
@ -41,39 +42,75 @@ impl From<PositionEncoding> for ruff_source_file::PositionEncoding {
/// A unique document ID, derived from a URL passed as part of an LSP request.
/// This document ID can point to either be a standalone Python file, a full notebook, or a cell within a notebook.
#[derive(Clone, Debug)]
pub(crate) enum DocumentKey {
Notebook(AnySystemPath),
NotebookCell {
cell_url: Url,
notebook_path: AnySystemPath,
},
Text(AnySystemPath),
///
/// The `DocumentKey` is very similar to `AnySystemPath`. The important distinction is that
/// ty doesn't know about individual notebook cells, instead, ty operates on full notebook documents.
/// ty also doesn't support resolving settings per cell, instead, settings are resolved per file or notebook.
///
/// Thus, the motivation of `DocumentKey` is to prevent accidental use of Cell keys for operations
/// that expect to work on a file path level. That's what [`DocumentHandle::to_file_path`]
/// is for, it returns a file path for any document, taking into account that these methods should
/// return the notebook for cell documents and notebooks.
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub(super) enum DocumentKey {
/// A URI using the `file` schema and maps to a valid path.
File(SystemPathBuf),
/// Any other URI.
///
/// Used for Notebook-cells, URI's with non-`file` schemes, or invalid `file` URI's.
Opaque(String),
}
impl DocumentKey {
/// Returns the file path associated with the key.
pub(crate) fn path(&self) -> &AnySystemPath {
match self {
DocumentKey::Notebook(path) | DocumentKey::Text(path) => path,
DocumentKey::NotebookCell { notebook_path, .. } => notebook_path,
/// Converts the given [`Url`] to an [`DocumentKey`].
///
/// If the URL scheme is `file`, then the path is converted to a [`SystemPathBuf`] unless
/// the url isn't a valid file path.
///
/// In all other cases, the URL is kept as an opaque identifier ([`Self::Opaque`]).
pub(crate) fn from_url(url: &Url) -> Self {
if url.scheme() == "file" {
if let Ok(path) = url.to_file_path() {
Self::File(SystemPathBuf::from_path_buf(path).expect("URL to be valid UTF-8"))
} else {
tracing::warn!(
"Treating `file:` url `{url}` as opaque URL as it isn't a valid file path"
);
Self::Opaque(url.to_string())
}
} else {
Self::Opaque(url.to_string())
}
}
pub(crate) fn from_path(path: AnySystemPath) -> Self {
// For text documents, we assume it's a text document unless it's a notebook file.
match path.extension() {
Some("ipynb") => Self::Notebook(path),
_ => Self::Text(path),
pub(crate) fn as_opaque(&self) -> Option<&str> {
match self {
Self::Opaque(uri) => Some(uri),
Self::File(_) => None,
}
}
/// Returns the URL for this document key. For notebook cells, returns the cell URL.
/// For other document types, converts the path to a URL.
pub(crate) fn to_url(&self) -> Option<Url> {
/// Returns the corresponding [`AnySystemPath`] for this document key.
///
/// Note, calling this method on a `DocumentKey::Opaque` representing a cell document
/// will return a `SystemVirtualPath` corresponding to the cell URI but not the notebook file path.
/// That's most likely not what you want.
pub(super) fn to_file_path(&self) -> AnySystemPath {
match self {
DocumentKey::NotebookCell { cell_url, .. } => Some(cell_url.clone()),
DocumentKey::Notebook(path) | DocumentKey::Text(path) => path.to_url(),
Self::File(path) => AnySystemPath::System(path.clone()),
Self::Opaque(uri) => {
AnySystemPath::SystemVirtual(SystemVirtualPath::new(uri).to_path_buf())
}
}
}
}
impl From<AnySystemPath> for DocumentKey {
fn from(value: AnySystemPath) -> Self {
match value {
AnySystemPath::System(system_path) => Self::File(system_path),
AnySystemPath::SystemVirtual(virtual_path) => Self::Opaque(virtual_path.to_string()),
}
}
}
@ -81,11 +118,8 @@ impl DocumentKey {
impl std::fmt::Display for DocumentKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NotebookCell { cell_url, .. } => cell_url.fmt(f),
Self::Notebook(path) | Self::Text(path) => match path {
AnySystemPath::System(system_path) => system_path.fmt(f),
AnySystemPath::SystemVirtual(virtual_path) => virtual_path.fmt(f),
},
Self::File(path) => path.fmt(f),
Self::Opaque(uri) => uri.fmt(f),
}
}
}

100
crates/ty_server/src/document/notebook.rs
View File

@ -3,9 +3,8 @@ use lsp_types::NotebookCellKind;
use ruff_notebook::CellMetadata;
use rustc_hash::{FxBuildHasher, FxHashMap};
use crate::{PositionEncoding, TextDocument};
use super::DocumentVersion;
use crate::{PositionEncoding, TextDocument};
pub(super) type CellId = usize;
@ -13,16 +12,25 @@ pub(super) type CellId = usize;
/// contents are internally represented by [`TextDocument`]s.
#[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<lsp_types::Url, CellId>,
cell_index: FxHashMap<String, CellId>,
}
/// A single cell within a notebook, which has text contents represented as a `TextDocument`.
#[derive(Clone, Debug)]
struct NotebookCell {
/// The URL uniquely identifying the cell.
///
/// > Cell text documents have a URI, but servers should not rely on any
/// > format for this URI, since it is up to the client on how it will
/// > create these URIs. The URIs must be unique across ALL notebook
/// > cells and can therefore be used to uniquely identify a notebook cell
/// > or the cell’s text document.
/// > <https://microsoft.github.io/language-server-protocol/specifications/lsp/3.18/specification/#notebookDocument_synchronization>
url: lsp_types::Url,
kind: NotebookCellKind,
document: TextDocument,
@ -30,32 +38,45 @@ struct NotebookCell {
impl NotebookDocument {
pub fn new(
version: DocumentVersion,
url: lsp_types::Url,
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 cell_contents: FxHashMap<_, _> = cell_documents
.into_iter()
.map(|document| (document.uri, document.text))
.collect();
let mut cells: Vec<_> = cells.into_iter().map(NotebookCell::empty).collect();
let cells: Vec<_> = cells
.into_iter()
.map(|cell| {
let contents = cell_contents.remove(&cell.document).unwrap_or_default();
NotebookCell::new(cell, contents, version)
})
.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);
}
Ok(Self {
version,
cell_index: Self::make_cell_index(cells.as_slice()),
metadata: serde_json::from_value(serde_json::Value::Object(metadata))?,
url,
version: notebook_version,
cell_index,
cells,
metadata: serde_json::from_value(serde_json::Value::Object(metadata))?,
})
}
pub(crate) fn url(&self) -> &lsp_types::Url {
&self.url
}
/// 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 {
@ -127,7 +148,7 @@ impl NotebookDocument {
// 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);
self.cell_index.remove(cell.url.as_str());
deleted_cells.insert(cell.url, cell.document);
}
}
@ -150,7 +171,7 @@ impl NotebookDocument {
// 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.clone(), index);
self.cell_index.insert(cell.url.to_string(), index);
}
// Finally, update the text document that represents the cell with the actual
@ -158,8 +179,9 @@ impl NotebookDocument {
// `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) {
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,
);
@ -170,7 +192,7 @@ impl NotebookDocument {
if let Some(cell_data) = data {
for cell in cell_data {
if let Some(existing_cell) = self.cell_by_uri_mut(&cell.document) {
if let Some(existing_cell) = self.cell_by_uri_mut(cell.document.as_str()) {
existing_cell.kind = cell.kind;
}
}
@ -178,7 +200,7 @@ impl NotebookDocument {
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) {
if let Some(cell) = self.cell_by_uri_mut(content_change.document.uri.as_str()) {
cell.document
.apply_changes(content_change.changes, version, encoding);
}
@ -204,7 +226,8 @@ impl NotebookDocument {
}
/// Get the text document representing the contents of a cell by the cell URI.
pub(crate) fn cell_document_by_uri(&self, uri: &lsp_types::Url) -> Option<&TextDocument> {
#[expect(unused)]
pub(crate) fn cell_document_by_uri(&self, uri: &str) -> Option<&TextDocument> {
self.cells
.get(*self.cell_index.get(uri)?)
.map(|cell| &cell.document)
@ -215,29 +238,41 @@ impl NotebookDocument {
self.cells.iter().map(|cell| &cell.url)
}
fn cell_by_uri_mut(&mut self, uri: &lsp_types::Url) -> Option<&mut NotebookCell> {
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<lsp_types::Url, CellId> {
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.clone(), i);
index.insert(cell.url.to_string(), i);
}
index
}
}
impl NotebookCell {
pub(crate) fn empty(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,
document: TextDocument::new(contents, version),
}
}
}
@ -294,7 +329,14 @@ mod tests {
}
}
NotebookDocument::new(0, cells, serde_json::Map::default(), cell_documents).unwrap()
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

15
crates/ty_server/src/document/text_document.rs
View File

@ -1,4 +1,4 @@
use lsp_types::TextDocumentContentChangeEvent;
use lsp_types::{TextDocumentContentChangeEvent, Url};
use ruff_source_file::LineIndex;
use crate::PositionEncoding;
@ -11,6 +11,9 @@ pub(crate) type DocumentVersion = i32;
/// with changes made by the user, including unsaved changes.
#[derive(Debug, Clone)]
pub struct TextDocument {
/// The URL as sent by the client
url: Url,
/// The string contents of the document.
contents: String,
/// A computed line index for the document. This should always reflect
@ -40,9 +43,10 @@ impl From<&str> for LanguageId {
}
impl TextDocument {
pub fn new(contents: String, version: DocumentVersion) -> Self {
pub fn new(url: Url, contents: String, version: DocumentVersion) -> Self {
let index = LineIndex::from_source_text(&contents);
Self {
url,
contents,
index,
version,
@ -60,6 +64,10 @@ impl TextDocument {
self.contents
}
pub(crate) fn url(&self) -> &Url {
&self.url
}
pub fn contents(&self) -> &str {
&self.contents
}
@ -154,11 +162,12 @@ impl TextDocument {
#[cfg(test)]
mod tests {
use crate::{PositionEncoding, TextDocument};
use lsp_types::{Position, TextDocumentContentChangeEvent};
use lsp_types::{Position, TextDocumentContentChangeEvent, Url};
#[test]
fn redo_edit() {
let mut document = TextDocument::new(
Url::parse("file:///test").unwrap(),
r#""""
测试comment
一些测试内容

2
crates/ty_server/src/lib.rs
View File

@ -8,7 +8,7 @@ pub use crate::logging::{LogLevel, init_logging};
pub use crate::server::{PartialWorkspaceProgress, PartialWorkspaceProgressParams, Server};
pub use crate::session::{ClientOptions, DiagnosticMode};
pub use document::{NotebookDocument, PositionEncoding, TextDocument};
pub(crate) use session::{DocumentQuery, Session};
pub(crate) use session::Session;
mod capabilities;
mod document;

23
crates/ty_server/src/server/api.rs
View File

@ -1,6 +1,5 @@
use crate::server::schedule::Task;
use crate::session::Session;
use crate::system::AnySystemPath;
use anyhow::anyhow;
use lsp_server as server;
use lsp_server::RequestId;
@ -208,7 +207,7 @@ where
// SAFETY: The `snapshot` is safe to move across the unwind boundary because it is not used
// after unwinding.
let snapshot = AssertUnwindSafe(session.take_session_snapshot());
let snapshot = AssertUnwindSafe(session.snapshot_session());
Box::new(move |client| {
let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();
@ -253,10 +252,10 @@ where
.cancellation_token(&id)
.expect("request should have been tested for cancellation before scheduling");
let url = R::document_url(&params).into_owned();
let url = R::document_url(&params);
let Ok(path) = AnySystemPath::try_from_url(&url) else {
let reason = format!("URL `{url}` isn't a valid system path");
let Ok(document) = session.snapshot_document(&url) else {
let reason = format!("Document {url} is not open in the session");
tracing::warn!(
"Ignoring request id={id} method={} because {reason}",
R::METHOD
@ -274,8 +273,8 @@ where
});
};
let path = document.to_file_path();
let db = session.project_db(&path).clone();
let snapshot = session.take_document_snapshot(url);
Box::new(move |client| {
let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();
@ -294,7 +293,7 @@ where
}
if let Err(error) = ruff_db::panic::catch_unwind(|| {
R::handle_request(&id, &db, snapshot, client, params);
R::handle_request(&id, &db, document, client, params);
}) {
panic_response::<R>(&id, client, &error, retry);
}
@ -371,7 +370,15 @@ where
let (id, params) = cast_notification::<N>(req)?;
Ok(Task::background(schedule, move |session: &Session| {
let url = N::document_url(&params);
let snapshot = session.take_document_snapshot((*url).clone());
let Ok(snapshot) = session.snapshot_document(&url) else {
let reason = format!("Document {url} is not open in the session");
tracing::warn!(
"Ignoring notification id={id} method={} because {reason}",
N::METHOD
);
return Box::new(|_| {});
};
Box::new(move |client| {
let _span = tracing::debug_span!("notification", method = N::METHOD).entered();

52
crates/ty_server/src/server/api/diagnostics.rs
View File

@ -13,16 +13,16 @@ use ruff_db::source::{line_index, source_text};
use ruff_db::system::SystemPathBuf;
use ty_project::{Db, ProjectDatabase};
use crate::document::{DocumentKey, FileRangeExt, ToRangeExt};
use crate::document::{FileRangeExt, ToRangeExt};
use crate::session::DocumentSnapshot;
use crate::session::client::Client;
use crate::system::{AnySystemPath, file_to_url};
use crate::{DocumentQuery, PositionEncoding, Session};
use crate::{NotebookDocument, PositionEncoding, Session};
pub(super) struct Diagnostics<'a> {
items: Vec<ruff_db::diagnostic::Diagnostic>,
encoding: PositionEncoding,
document: &'a DocumentQuery,
notebook: Option<&'a NotebookDocument>,
}
impl Diagnostics<'_> {
@ -53,7 +53,7 @@ impl Diagnostics<'_> {
}
pub(super) fn to_lsp_diagnostics(&self, db: &ProjectDatabase) -> LspDiagnostics {
if let Some(notebook) = self.document.as_notebook() {
if let Some(notebook) = self.notebook {
let mut cell_diagnostics: FxHashMap<Url, Vec<Diagnostic>> = FxHashMap::default();
// Populates all relevant URLs with an empty diagnostic list. This ensures that documents
@ -115,23 +115,18 @@ impl LspDiagnostics {
}
}
/// Clears the diagnostics for the document identified by `key`.
/// 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, key: &DocumentKey, client: &Client) {
pub(super) fn clear_diagnostics(session: &Session, uri: &lsp_types::Url, client: &Client) {
if session.client_capabilities().supports_pull_diagnostics() {
return;
}
let Some(uri) = key.to_url() else {
// If we can't convert to URL, we can't clear diagnostics
return;
};
client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
uri,
uri: uri.clone(),
diagnostics: vec![],
version: None,
});
@ -143,18 +138,12 @@ pub(super) fn clear_diagnostics(session: &Session, key: &DocumentKey, client: &C
/// This function is a no-op if the client supports pull diagnostics.
///
/// [publish diagnostics notification]: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_publishDiagnostics
pub(super) fn publish_diagnostics(session: &Session, key: &DocumentKey, client: &Client) {
pub(super) fn publish_diagnostics(session: &Session, url: &lsp_types::Url, client: &Client) {
if session.client_capabilities().supports_pull_diagnostics() {
return;
}
let Some(url) = key.to_url() else {
return;
};
let snapshot = session.take_document_snapshot(url.clone());
let document = match snapshot.document() {
let snapshot = match session.snapshot_document(url) {
Ok(document) => document,
Err(err) => {
tracing::debug!("Failed to resolve document for URL `{}`: {}", url, err);
@ -162,7 +151,7 @@ pub(super) fn publish_diagnostics(session: &Session, key: &DocumentKey, client:
}
};
let db = session.project_db(key.path());
let db = session.project_db(&snapshot.to_file_path());
let Some(diagnostics) = compute_diagnostics(db, &snapshot) else {
return;
@ -173,13 +162,13 @@ pub(super) fn publish_diagnostics(session: &Session, key: &DocumentKey, client:
client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
uri,
diagnostics,
version: Some(document.version()),
version: Some(snapshot.document().version()),
});
};
match diagnostics.to_lsp_diagnostics(db) {
LspDiagnostics::TextDocument(diagnostics) => {
publish_diagnostics_notification(url, diagnostics);
publish_diagnostics_notification(url.clone(), diagnostics);
}
LspDiagnostics::NotebookDocument(cell_diagnostics) => {
for (cell_url, diagnostics) in cell_diagnostics {
@ -264,16 +253,11 @@ pub(super) fn compute_diagnostics<'a>(
db: &ProjectDatabase,
snapshot: &'a DocumentSnapshot,
) -> Option<Diagnostics<'a>> {
let document = match snapshot.document() {
Ok(document) => document,
Err(err) => {
tracing::info!("Failed to resolve document for snapshot: {}", err);
return None;
}
};
let Some(file) = document.file(db) else {
tracing::info!("No file found for snapshot for `{}`", document.file_path());
let Some(file) = snapshot.to_file(db) else {
tracing::info!(
"No file found for snapshot for `{}`",
snapshot.to_file_path()
);
return None;
};
@ -282,7 +266,7 @@ pub(super) fn compute_diagnostics<'a>(
Some(Diagnostics {
items: diagnostics,
encoding: snapshot.encoding(),
document,
notebook: snapshot.notebook(),
})
}

23
crates/ty_server/src/server/api/notifications/did_change.rs
View File

@ -28,19 +28,16 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
content_changes,
} = params;
let key = match session.key_from_url(uri) {
Ok(key) => key,
Err(uri) => {
tracing::debug!("Failed to create document key from URI: {}", uri);
return Ok(());
}
};
session
.update_text_document(&key, content_changes, version)
let document = session
.document_handle(&uri)
.with_failure_code(ErrorCode::InternalError)?;
let changes = match key.path() {
document
.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())]
}
@ -49,9 +46,9 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
}
};
session.apply_changes(key.path(), changes);
session.apply_changes(&path, changes);
publish_diagnostics(session, &key, client);
publish_diagnostics(session, document.url(), client);
Ok(())
}

17
crates/ty_server/src/server/api/notifications/did_change_watched_files.rs
View File

@ -1,3 +1,4 @@
use crate::document::DocumentKey;
use crate::server::Result;
use crate::server::api::diagnostics::{publish_diagnostics, publish_settings_diagnostics};
use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
@ -25,16 +26,8 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
let mut events_by_db: FxHashMap<_, Vec<ChangeEvent>> = FxHashMap::default();
for change in params.changes {
let path = match AnySystemPath::try_from_url(&change.uri) {
Ok(path) => path,
Err(err) => {
tracing::warn!(
"Failed to convert URI '{}` to system path: {err:?}",
change.uri
);
continue;
}
};
let key = DocumentKey::from_url(&change.uri);
let path = key.to_file_path();
let system_path = match path {
AnySystemPath::System(system) => system,
@ -99,8 +92,8 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
|_, ()| {},
);
} else {
for key in session.text_document_keys() {
publish_diagnostics(session, &key, client);
for key in session.text_document_handles() {
publish_diagnostics(session, key.url(), client);
}
}
// TODO: always publish diagnostics for notebook files (since they don't use pull diagnostics)

28
crates/ty_server/src/server/api/notifications/did_close.rs
View File

@ -27,22 +27,20 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
text_document: TextDocumentIdentifier { uri },
} = params;
let key = match session.key_from_url(uri) {
Ok(key) => key,
Err(uri) => {
tracing::debug!("Failed to create document key from URI: {}", uri);
return Ok(());
}
};
session
.close_document(&key)
let document = session
.document_handle(&uri)
.with_failure_code(ErrorCode::InternalError)?;
let path = key.path();
let db = session.project_db_mut(path);
let path = document.to_file_path().into_owned();
let url = document.url().clone();
match path {
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);
@ -65,7 +63,7 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
.diagnostic_mode()
.is_open_files_only()
{
clear_diagnostics(session, &key, client);
clear_diagnostics(session, &url, client);
}
}
AnySystemPath::SystemVirtual(virtual_path) => {
@ -78,7 +76,7 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
// 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, &key, client);
clear_diagnostics(session, &url, client);
}
}

22
crates/ty_server/src/server/api/notifications/did_close_notebook.rs
View File

@ -26,21 +26,19 @@ impl SyncNotificationHandler for DidCloseNotebookHandler {
..
} = params;
let key = match session.key_from_url(uri) {
Ok(key) => key,
Err(uri) => {
tracing::debug!("Failed to create document key from URI: {}", uri);
return Ok(());
}
};
session
.close_document(&key)
let document = session
.document_handle(&uri)
.with_failure_code(lsp_server::ErrorCode::InternalError)?;
if let AnySystemPath::SystemVirtual(virtual_path) = key.path() {
let path = document.to_file_path().into_owned();
document
.close(session)
.with_failure_code(lsp_server::ErrorCode::InternalError)?;
if let AnySystemPath::SystemVirtual(virtual_path) = &path {
session.apply_changes(
key.path(),
&path,
vec![ChangeEvent::DeletedVirtual(virtual_path.clone())],
);
}

27
crates/ty_server/src/server/api/notifications/did_open.rs
View File

@ -35,30 +35,23 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
},
} = params;
let key = match session.key_from_url(uri) {
Ok(key) => key,
Err(uri) => {
tracing::debug!("Failed to create document key from URI: {}", uri);
return Ok(());
}
};
let document = session.open_text_document(
TextDocument::new(uri, text, version).with_language_id(&language_id),
);
let document = TextDocument::new(text, version).with_language_id(&language_id);
session.open_text_document(key.path(), document);
let path = key.path();
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);
let db = session.project_db(&path);
db.files()
.try_system(db, system_path)
.is_none_or(|file| !file.exists(db))
});
match path {
match &*path {
AnySystemPath::System(system_path) => {
let event = if is_maybe_new_system_file {
ChangeEvent::Created {
@ -68,22 +61,22 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
} else {
ChangeEvent::Opened(system_path.clone())
};
session.apply_changes(path, vec![event]);
session.apply_changes(&path, vec![event]);
let db = session.project_db_mut(path);
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 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, &key, client);
publish_diagnostics(session, document.url(), client);
Ok(())
}

23
crates/ty_server/src/server/api/notifications/did_open_notebook.rs
View File

@ -25,20 +25,27 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
_client: &Client,
params: DidOpenNotebookDocumentParams,
) -> Result<()> {
let Ok(path) = AnySystemPath::try_from_url(&params.notebook_document.uri) else {
return Ok(());
};
let lsp_types::NotebookDocument {
version,
cells,
metadata,
uri: notebook_uri,
..
} = params.notebook_document;
let notebook = NotebookDocument::new(
params.notebook_document.version,
params.notebook_document.cells,
params.notebook_document.metadata.unwrap_or_default(),
notebook_uri,
version,
cells,
metadata.unwrap_or_default(),
params.cell_text_documents,
)
.with_failure_code(ErrorCode::InternalError)?;
session.open_notebook_document(&path, notebook);
match &path {
let document = session.open_notebook_document(notebook);
let path = document.to_file_path();
match &*path {
AnySystemPath::System(system_path) => {
session.apply_changes(&path, vec![ChangeEvent::Opened(system_path.clone())]);
}

2
crates/ty_server/src/server/api/requests/completion.rs
View File

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

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

@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for DocumentHighlightRequestHandler {
return Ok(None);
}
let Some(file) = snapshot.file(db) else {
let Some(file) = snapshot.to_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.file(db) else {
let Some(file) = snapshot.to_file(db) else {
return Ok(None);
};

2
crates/ty_server/src/server/api/requests/execute_command.rs
View File

@ -52,7 +52,7 @@ fn debug_information(session: &Session) -> crate::Result<String> {
writeln!(
buffer,
"Open text documents: {}",
session.text_document_keys().count()
session.text_document_handles().count()
)?;
writeln!(buffer)?;

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

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

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

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

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

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

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

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

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

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

2
crates/ty_server/src/server/api/requests/inlay_hints.rs
View File

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

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

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

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

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

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

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

2
crates/ty_server/src/server/api/requests/semantic_tokens.rs
View File

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

2
crates/ty_server/src/server/api/requests/semantic_tokens_range.rs
View File

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

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

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

31
crates/ty_server/src/server/api/requests/workspace_diagnostic.rs
View File

@ -1,4 +1,5 @@
use crate::PositionEncoding;
use crate::document::DocumentKey;
use crate::server::api::diagnostics::{Diagnostics, to_lsp_diagnostic};
use crate::server::api::traits::{
BackgroundRequestHandler, RequestHandler, RetriableRequestHandler,
@ -8,7 +9,7 @@ use crate::server::{Action, Result};
use crate::session::client::Client;
use crate::session::index::Index;
use crate::session::{SessionSnapshot, SuspendedWorkspaceDiagnosticRequest};
use crate::system::{AnySystemPath, file_to_url};
use crate::system::file_to_url;
use lsp_server::RequestId;
use lsp_types::request::WorkspaceDiagnosticRequest;
use lsp_types::{
@ -317,7 +318,7 @@ struct ResponseWriter<'a> {
// It's important that we use `AnySystemPath` over `Url` here because
// `file_to_url` isn't guaranteed to return the exact same URL as the one provided
// by the client.
previous_result_ids: FxHashMap<AnySystemPath, (Url, String)>,
previous_result_ids: FxHashMap<DocumentKey, (Url, String)>,
}
impl<'a> ResponseWriter<'a> {
@ -346,12 +347,7 @@ impl<'a> ResponseWriter<'a> {
let previous_result_ids = previous_result_ids
.into_iter()
.filter_map(|prev| {
Some((
AnySystemPath::try_from_url(&prev.uri).ok()?,
(prev.uri, prev.value),
))
})
.map(|prev| (DocumentKey::from_url(&prev.uri), (prev.uri, prev.value)))
.collect();
Self {
@ -367,20 +363,16 @@ impl<'a> ResponseWriter<'a> {
tracing::debug!("Failed to convert file path to URL at {}", file.path(db));
return;
};
let key = DocumentKey::from_url(&url);
let version = self
.index
.key_from_url(url.clone())
.ok()
.and_then(|key| self.index.make_document_ref(key).ok())
.map(|doc| i64::from(doc.version()));
.document_handle(&url)
.map(|doc| i64::from(doc.version()))
.ok();
let result_id = Diagnostics::result_id_from_hash(diagnostics);
let previous_result_id = AnySystemPath::try_from_url(&url)
.ok()
.and_then(|path| self.previous_result_ids.remove(&path))
.map(|(_url, id)| id);
let previous_result_id = self.previous_result_ids.remove(&key).map(|(_url, id)| id);
let report = match result_id {
Some(new_id) if Some(&new_id) == previous_result_id.as_ref() => {
@ -444,13 +436,12 @@ impl<'a> ResponseWriter<'a> {
// Handle files that had diagnostics in previous request but no longer have any
// Any remaining entries in previous_results are files that were fixed
for (previous_url, previous_result_id) in self.previous_result_ids.into_values() {
for (key, (previous_url, previous_result_id)) in self.previous_result_ids {
// This file had diagnostics before but doesn't now, so we need to report it as having no diagnostics
let version = self
.index
.key_from_url(previous_url.clone())
.document(&key)
.ok()
.and_then(|key| self.index.make_document_ref(key).ok())
.map(|doc| i64::from(doc.version()));
let new_result_id = Diagnostics::result_id_from_hash(&[]);

230
crates/ty_server/src/session.rs
View File

@ -1,7 +1,7 @@
//! Data model, state management, and configuration resolution.
use anyhow::{Context, anyhow};
use index::DocumentQueryError;
use index::DocumentError;
use lsp_server::{Message, RequestId};
use lsp_types::notification::{DidChangeWatchedFiles, Exit, Notification};
use lsp_types::request::{
@ -15,8 +15,9 @@ use lsp_types::{
};
use options::GlobalOptions;
use ruff_db::Db;
use ruff_db::files::File;
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;
@ -26,7 +27,6 @@ use ty_project::metadata::Options;
use ty_project::watch::ChangeEvent;
use ty_project::{ChangeResult, CheckMode, Db as _, ProjectDatabase, ProjectMetadata};
pub(crate) use self::index::DocumentQuery;
pub(crate) use self::options::InitializationOptions;
pub use self::options::{ClientOptions, DiagnosticMode};
pub(crate) use self::settings::{GlobalSettings, WorkspaceSettings};
@ -439,13 +439,6 @@ impl Session {
self.projects.values_mut().chain(default_project)
}
/// Returns the [`DocumentKey`] for the given URL.
///
/// Refer to [`Index::key_from_url`] for more details.
pub(crate) fn key_from_url(&self, url: Url) -> Result<DocumentKey, Url> {
self.index().key_from_url(url)
}
pub(crate) fn initialize_workspaces(
&mut self,
workspace_settings: Vec<(Url, ClientOptions)>,
@ -819,25 +812,34 @@ impl Session {
}
/// Creates a document snapshot with the URL referencing the document to snapshot.
pub(crate) fn take_document_snapshot(&self, url: Url) -> DocumentSnapshot {
let key = self
.key_from_url(url)
.map_err(DocumentQueryError::InvalidUrl);
DocumentSnapshot {
pub(crate) fn snapshot_document(&self, url: &Url) -> Result<DocumentSnapshot, DocumentError> {
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: key
.as_ref()
.ok()
.and_then(|key| self.workspaces.settings_for_path(key.path().as_system()?))
workspace_settings: document_handle
.to_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_query_result: key.and_then(|key| self.index().make_document_ref(key)),
}
document: document_handle,
notebook,
})
}
/// Creates a snapshot of the current state of the [`Session`].
pub(crate) fn take_session_snapshot(&self) -> SessionSnapshot {
pub(crate) fn snapshot_session(&self) -> SessionSnapshot {
SessionSnapshot {
projects: self
.projects
@ -855,56 +857,49 @@ impl Session {
}
/// Iterates over the document keys for all open text documents.
pub(super) fn text_document_keys(&self) -> impl Iterator<Item = DocumentKey> + '_ {
pub(super) fn text_document_handles(&self) -> impl Iterator<Item = DocumentHandle> + '_ {
self.index()
.text_document_paths()
.map(|path| DocumentKey::Text(path.clone()))
.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,
})
}
/// Returns a handle to the document specified by its URL.
///
/// # Errors
///
/// If the document is not found.
pub(crate) fn document_handle(
&self,
url: &lsp_types::Url,
) -> Result<DocumentHandle, DocumentError> {
self.index().document_handle(url)
}
/// Registers a notebook document at the provided `path`.
/// If a document is already open here, it will be overwritten.
pub(crate) fn open_notebook_document(
&mut self,
path: &AnySystemPath,
document: NotebookDocument,
) {
self.index_mut().open_notebook_document(path, document);
///
/// 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();
handle
}
/// Registers a text document at the provided `path`.
/// If a document is already open here, it will be overwritten.
pub(crate) fn open_text_document(&mut self, path: &AnySystemPath, document: TextDocument) {
self.index_mut().open_text_document(path, document);
self.bump_revision();
}
/// Updates a text document at the associated `key`.
///
/// The document key must point to a text document, or this will throw an error.
pub(crate) fn update_text_document(
&mut self,
key: &DocumentKey,
content_changes: Vec<TextDocumentContentChangeEvent>,
new_version: DocumentVersion,
) -> crate::Result<()> {
let position_encoding = self.position_encoding;
self.index_mut().update_text_document(
key,
content_changes,
new_version,
position_encoding,
)?;
self.bump_revision();
Ok(())
}
/// 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);
/// De-registers a document, specified by its key.
/// Calling this multiple times for the same document is a logic error.
pub(crate) fn close_document(&mut self, key: &DocumentKey) -> crate::Result<()> {
self.index_mut().close_document(key)?;
self.bump_revision();
Ok(())
handle
}
/// Returns a reference to the index.
@ -1003,7 +998,8 @@ pub(crate) struct DocumentSnapshot {
global_settings: Arc<GlobalSettings>,
workspace_settings: Arc<WorkspaceSettings>,
position_encoding: PositionEncoding,
document_query_result: Result<DocumentQuery, DocumentQueryError>,
document: DocumentHandle,
notebook: Option<Arc<NotebookDocument>>,
}
impl DocumentSnapshot {
@ -1028,27 +1024,28 @@ impl DocumentSnapshot {
}
/// Returns the result of the document query for this snapshot.
pub(crate) fn document(&self) -> Result<&DocumentQuery, &DocumentQueryError> {
self.document_query_result.as_ref()
pub(crate) fn document(&self) -> &DocumentHandle {
&self.document
}
pub(crate) fn file(&self, db: &dyn Db) -> Option<File> {
let document = match self.document() {
Ok(document) => document,
Err(err) => {
tracing::debug!("Failed to resolve file: {}", err);
return None;
}
};
let file = document.file(db);
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);
if file.is_none() {
tracing::debug!(
"Failed to resolve file: file not found for path `{}`",
document.file_path()
"Failed to resolve file: file not found for `{}`",
self.document.url()
);
}
file
}
pub(crate) fn to_file_path(&self) -> Cow<'_, AnySystemPath> {
self.document.to_file_path()
}
}
/// An immutable snapshot of the current state of [`Session`].
@ -1320,3 +1317,90 @@ impl SuspendedWorkspaceDiagnosticRequest {
None
}
}
/// A handle to a document stored within [`Index`].
///
/// Allows identifying the document within the index but it also carries the URL used by the
/// client to reference the document as well as the version of the document.
///
/// 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,
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>,
version: DocumentVersion,
}
impl DocumentHandle {
pub(crate) const fn version(&self) -> DocumentVersion {
self.version
}
/// The URL as used by the client to reference this document.
pub(crate) fn url(&self) -> &lsp_types::Url {
&self.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())
}
}
/// Returns the salsa interned [`File`] for the document selected by this query.
///
/// 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() {
AnySystemPath::System(path) => system_path_to_file(db, path).ok(),
AnySystemPath::SystemVirtual(virtual_path) => db
.files()
.try_virtual_file(virtual_path)
.map(|virtual_file| virtual_file.file()),
}
}
pub(crate) fn update_text_document(
&self,
session: &mut Session,
content_changes: Vec<TextDocumentContentChangeEvent>,
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 Some(document) = document_mut.as_text_mut() else {
anyhow::bail!("Text document path does not point to a text document");
};
if content_changes.is_empty() {
document.update_version(new_version);
return Ok(());
}
document.apply_changes(content_changes, new_version, position_encoding);
Ok(())
}
/// 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)?;
session.bump_revision();
Ok(())
}
}

314
crates/ty_server/src/session/index.rs
View File

@ -1,24 +1,24 @@
use std::sync::Arc;
use lsp_types::Url;
use ruff_db::Db;
use ruff_db::files::{File, system_path_to_file};
use rustc_hash::FxHashMap;
use crate::document::DocumentKey;
use crate::session::DocumentHandle;
use crate::{
PositionEncoding, TextDocument,
document::{DocumentKey, DocumentVersion, NotebookDocument},
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<AnySystemPath, DocumentController>,
documents: FxHashMap<DocumentKey, Document>,
/// Maps opaque cell URLs to a notebook path (document)
notebook_cells: FxHashMap<Url, AnySystemPath>,
notebook_cells: FxHashMap<String, AnySystemPath>,
}
impl Index {
@ -29,68 +29,55 @@ impl Index {
}
}
pub(super) fn text_document_paths(&self) -> impl Iterator<Item = &AnySystemPath> + '_ {
self.documents
.iter()
.filter_map(|(path, doc)| doc.as_text().and(Some(path)))
pub(super) fn text_documents(
&self,
) -> impl Iterator<Item = (&DocumentKey, &TextDocument)> + '_ {
self.documents.iter().filter_map(|(key, doc)| {
let text_document = doc.as_text()?;
Some((key, text_document))
})
}
pub(crate) fn document_handle(
&self,
url: &lsp_types::Url,
) -> Result<DocumentHandle, DocumentError> {
let key = DocumentKey::from_url(url);
let Some(document) = self.documents.get(&key) else {
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(),
})
}
#[expect(dead_code)]
pub(super) fn notebook_document_paths(&self) -> impl Iterator<Item = &AnySystemPath> + '_ {
pub(super) fn notebook_document_keys(&self) -> impl Iterator<Item = &DocumentKey> + '_ {
self.documents
.iter()
.filter(|(_, doc)| doc.as_notebook().is_some())
.map(|(path, _)| path)
}
pub(super) fn update_text_document(
&mut self,
key: &DocumentKey,
content_changes: Vec<lsp_types::TextDocumentContentChangeEvent>,
new_version: DocumentVersion,
encoding: PositionEncoding,
) -> crate::Result<()> {
let controller = self.document_controller_for_key(key)?;
let Some(document) = controller.as_text_mut() else {
anyhow::bail!("Text document path does not point to a text document");
};
if content_changes.is_empty() {
document.update_version(new_version);
return Ok(());
}
document.apply_changes(content_changes, new_version, encoding);
Ok(())
}
/// Returns the [`DocumentKey`] corresponding to the given URL.
///
/// It returns [`Err`] with the original URL if it cannot be converted to a [`AnySystemPath`].
pub(crate) fn key_from_url(&self, url: Url) -> Result<DocumentKey, Url> {
if let Some(notebook_path) = self.notebook_cells.get(&url) {
Ok(DocumentKey::NotebookCell {
cell_url: url,
notebook_path: notebook_path.clone(),
})
} else {
let path = AnySystemPath::try_from_url(&url).map_err(|()| url)?;
if path
.extension()
.is_some_and(|ext| ext.eq_ignore_ascii_case("ipynb"))
{
Ok(DocumentKey::Notebook(path))
} else {
Ok(DocumentKey::Text(path))
}
}
.map(|(key, _)| key)
}
#[expect(dead_code)]
pub(super) fn update_notebook_document(
&mut self,
key: &DocumentKey,
notebook_key: &DocumentKey,
cells: Option<lsp_types::NotebookDocumentCellChange>,
metadata: Option<serde_json::Map<String, serde_json::Value>>,
new_version: DocumentVersion,
@ -102,17 +89,16 @@ impl Index {
..
}) = cells.as_ref().and_then(|cells| cells.structure.as_ref())
{
let notebook_path = key.path().clone();
for opened_cell in did_open {
let cell_path = SystemVirtualPath::new(opened_cell.uri.as_str());
self.notebook_cells
.insert(opened_cell.uri.clone(), notebook_path.clone());
.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 controller = self.document_controller_for_key(key)?;
let Some(notebook) = controller.as_notebook_mut() else {
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");
};
@ -123,44 +109,64 @@ impl Index {
/// Create a document reference corresponding to the given document key.
///
/// Returns an error if the document is not found or if the path cannot be converted to a URL.
pub(crate) fn make_document_ref(
pub(crate) fn document(&self, key: &DocumentKey) -> Result<&Document, DocumentError> {
let Some(document) = self.documents.get(key) else {
return Err(DocumentError::NotFound(key.clone()));
};
Ok(document)
}
pub(crate) fn notebook_arc(
&self,
key: DocumentKey,
) -> Result<DocumentQuery, DocumentQueryError> {
let path = key.path();
let Some(controller) = self.documents.get(path) else {
return Err(DocumentQueryError::NotFound(key));
key: &DocumentKey,
) -> Result<Arc<NotebookDocument>, DocumentError> {
let Some(document) = self.documents.get(key) else {
return Err(DocumentError::NotFound(key.clone()));
};
// TODO: The `to_url` conversion shouldn't be an error because the paths themselves are
// constructed from the URLs but the `Index` APIs don't maintain this invariant.
let (cell_url, file_path) = match key {
DocumentKey::NotebookCell {
cell_url,
notebook_path,
} => (Some(cell_url), notebook_path),
DocumentKey::Notebook(path) | DocumentKey::Text(path) => (None, path),
};
Ok(controller.make_ref(cell_url, file_path))
if let Document::Notebook(notebook) = document {
Ok(notebook.clone())
} else {
Err(DocumentError::NotFound(key.clone()))
}
}
pub(super) fn open_text_document(&mut self, path: &AnySystemPath, document: TextDocument) {
self.documents
.insert(path.clone(), DocumentController::new_text(document));
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(),
};
self.documents.insert(key, Document::new_text(document));
handle
}
pub(super) fn open_notebook_document(
&mut self,
notebook_path: &AnySystemPath,
document: NotebookDocument,
) {
pub(super) fn open_notebook_document(&mut self, document: NotebookDocument) -> DocumentHandle {
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.clone(), notebook_path.clone());
.insert(cell_url.to_string(), notebook_key.to_file_path());
}
self.documents
.insert(notebook_key.clone(), Document::new_notebook(document));
DocumentHandle {
notebook_path: Some(notebook_key.to_file_path()),
key: notebook_key,
url,
version,
}
self.documents.insert(
notebook_path.clone(),
DocumentController::new_notebook(document),
);
}
pub(super) fn close_document(&mut self, key: &DocumentKey) -> crate::Result<()> {
@ -169,27 +175,23 @@ impl Index {
// 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::NotebookCell { cell_url, .. } = key {
if self.notebook_cells.remove(cell_url).is_none() {
tracing::warn!("Tried to remove a notebook cell that does not exist: {cell_url}");
}
return Ok(());
if let DocumentKey::Opaque(uri) = key {
self.notebook_cells.remove(uri);
}
let path = key.path();
let Some(_) = self.documents.remove(path) else {
let Some(_) = self.documents.remove(key) else {
anyhow::bail!("tried to close document that didn't exist at {key}")
};
Ok(())
}
fn document_controller_for_key(
pub(super) fn document_mut(
&mut self,
key: &DocumentKey,
) -> crate::Result<&mut DocumentController> {
let path = key.path();
let Some(controller) = self.documents.get_mut(path) else {
anyhow::bail!("Document controller not available at `{key}`");
) -> Result<&mut Document, DocumentError> {
let Some(controller) = self.documents.get_mut(key) else {
return Err(DocumentError::NotFound(key.clone()));
};
Ok(controller)
}
@ -197,31 +199,24 @@ impl Index {
/// A mutable handler to an underlying document.
#[derive(Debug)]
enum DocumentController {
pub(crate) enum Document {
Text(Arc<TextDocument>),
Notebook(Arc<NotebookDocument>),
}
impl DocumentController {
fn new_text(document: TextDocument) -> Self {
impl Document {
pub(super) fn new_text(document: TextDocument) -> Self {
Self::Text(Arc::new(document))
}
fn new_notebook(document: NotebookDocument) -> Self {
pub(super) fn new_notebook(document: NotebookDocument) -> Self {
Self::Notebook(Arc::new(document))
}
fn make_ref(&self, cell_url: Option<Url>, file_path: AnySystemPath) -> DocumentQuery {
match &self {
Self::Notebook(notebook) => DocumentQuery::Notebook {
cell_url,
file_path,
notebook: notebook.clone(),
},
Self::Text(document) => DocumentQuery::Text {
file_path,
document: document.clone(),
},
pub(crate) fn version(&self) -> DocumentVersion {
match self {
Self::Text(document) => document.version(),
Self::Notebook(notebook) => notebook.version(),
}
}
@ -254,85 +249,8 @@ impl DocumentController {
}
}
/// A read-only query to an open document.
///
/// This query can 'select' a text document, full notebook, or a specific notebook cell.
/// It also includes document settings.
#[derive(Debug, Clone)]
pub(crate) enum DocumentQuery {
Text {
file_path: AnySystemPath,
document: Arc<TextDocument>,
},
Notebook {
/// The selected notebook cell, if it exists.
cell_url: Option<Url>,
/// The path to the notebook.
file_path: AnySystemPath,
notebook: Arc<NotebookDocument>,
},
}
impl DocumentQuery {
/// Attempts to access the underlying notebook document that this query is selecting.
pub(crate) fn as_notebook(&self) -> Option<&NotebookDocument> {
match self {
Self::Notebook { notebook, .. } => Some(notebook),
Self::Text { .. } => None,
}
}
/// Get the version of document selected by this query.
pub(crate) fn version(&self) -> DocumentVersion {
match self {
Self::Text { document, .. } => document.version(),
Self::Notebook { notebook, .. } => notebook.version(),
}
}
/// Get the system path for the document selected by this query.
pub(crate) fn file_path(&self) -> &AnySystemPath {
match self {
Self::Text { file_path, .. } | Self::Notebook { file_path, .. } => file_path,
}
}
/// Attempt to access the single inner text document selected by the query.
/// If this query is selecting an entire notebook document, this will return `None`.
#[expect(dead_code)]
pub(crate) fn as_single_document(&self) -> Option<&TextDocument> {
match self {
Self::Text { document, .. } => Some(document),
Self::Notebook {
notebook,
cell_url: cell_uri,
..
} => cell_uri
.as_ref()
.and_then(|cell_uri| notebook.cell_document_by_uri(cell_uri)),
}
}
/// Returns the salsa interned [`File`] for the document selected by this query.
///
/// 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 file(&self, db: &dyn Db) -> Option<File> {
match self.file_path() {
AnySystemPath::System(path) => system_path_to_file(db, path).ok(),
AnySystemPath::SystemVirtual(virtual_path) => db
.files()
.try_virtual_file(virtual_path)
.map(|virtual_file| virtual_file.file()),
}
}
}
#[derive(Debug, Clone, thiserror::Error)]
pub(crate) enum DocumentQueryError {
#[error("invalid URL: {0}")]
InvalidUrl(Url),
pub(crate) enum DocumentError {
#[error("document not found for key: {0}")]
NotFound(DocumentKey),
}

74
crates/ty_server/src/system.rs
View File

@ -4,6 +4,8 @@ use std::fmt::Display;
use std::panic::RefUnwindSafe;
use std::sync::Arc;
use crate::document::DocumentKey;
use crate::session::index::{Document, Index};
use lsp_types::Url;
use ruff_db::file_revision::FileRevision;
use ruff_db::files::{File, FilePath};
@ -16,10 +18,6 @@ use ruff_notebook::{Notebook, NotebookError};
use ty_ide::cached_vendored_path;
use ty_python_semantic::Db;
use crate::DocumentQuery;
use crate::document::DocumentKey;
use crate::session::index::Index;
/// Returns a [`Url`] for the given [`File`].
pub(crate) fn file_to_url(db: &dyn Db, file: File) -> Option<Url> {
match file.path(db) {
@ -41,26 +39,6 @@ pub(crate) enum AnySystemPath {
}
impl AnySystemPath {
/// Converts the given [`Url`] to an [`AnySystemPath`].
///
/// If the URL scheme is `file`, then the path is converted to a [`SystemPathBuf`]. Otherwise, the
/// URL is converted to a [`SystemVirtualPathBuf`].
///
/// This fails in the following cases:
/// * The URL cannot be converted to a file path (refer to [`Url::to_file_path`]).
/// * If the URL is not a valid UTF-8 string.
pub(crate) fn try_from_url(url: &Url) -> std::result::Result<Self, ()> {
if url.scheme() == "file" {
Ok(AnySystemPath::System(
SystemPathBuf::from_path_buf(url.to_file_path()?).map_err(|_| ())?,
))
} else {
Ok(AnySystemPath::SystemVirtual(
SystemVirtualPath::new(url.as_str()).to_path_buf(),
))
}
}
pub(crate) const fn as_system(&self) -> Option<&SystemPathBuf> {
match self {
AnySystemPath::System(system_path_buf) => Some(system_path_buf),
@ -68,21 +46,11 @@ impl AnySystemPath {
}
}
/// Returns the extension of the path, if any.
pub(crate) fn extension(&self) -> Option<&str> {
#[expect(unused)]
pub(crate) const fn as_virtual(&self) -> Option<&SystemVirtualPath> {
match self {
AnySystemPath::System(system_path) => system_path.extension(),
AnySystemPath::SystemVirtual(virtual_path) => virtual_path.extension(),
}
}
/// Converts the path to a URL.
pub(crate) fn to_url(&self) -> Option<Url> {
match self {
AnySystemPath::System(system_path) => {
Url::from_file_path(system_path.as_std_path()).ok()
}
AnySystemPath::SystemVirtual(virtual_path) => Url::parse(virtual_path.as_str()).ok(),
AnySystemPath::SystemVirtual(path) => Some(path.as_path()),
AnySystemPath::System(_) => None,
}
}
}
@ -144,21 +112,17 @@ impl LSPSystem {
self.index.as_ref().unwrap()
}
fn make_document_ref(&self, path: AnySystemPath) -> Option<DocumentQuery> {
fn make_document_ref(&self, path: AnySystemPath) -> Option<&Document> {
let index = self.index();
let key = DocumentKey::from_path(path);
index.make_document_ref(key).ok()
index.document(&DocumentKey::from(path)).ok()
}
fn system_path_to_document_ref(&self, path: &SystemPath) -> Option<DocumentQuery> {
fn system_path_to_document_ref(&self, path: &SystemPath) -> Option<&Document> {
let any_path = AnySystemPath::System(path.to_path_buf());
self.make_document_ref(any_path)
}
fn system_virtual_path_to_document_ref(
&self,
path: &SystemVirtualPath,
) -> Option<DocumentQuery> {
fn system_virtual_path_to_document_ref(&self, path: &SystemVirtualPath) -> Option<&Document> {
let any_path = AnySystemPath::SystemVirtual(path.to_path_buf());
self.make_document_ref(any_path)
}
@ -170,7 +134,7 @@ impl System for LSPSystem {
if let Some(document) = document {
Ok(Metadata::new(
document_revision(&document),
document_revision(document),
None,
FileType::File,
))
@ -191,7 +155,7 @@ impl System for LSPSystem {
let document = self.system_path_to_document_ref(path);
match document {
Some(DocumentQuery::Text { document, .. }) => Ok(document.contents().to_string()),
Some(Document::Text(document)) => Ok(document.contents().to_string()),
_ => self.native_system.read_to_string(path),
}
}
@ -200,10 +164,8 @@ impl System for LSPSystem {
let document = self.system_path_to_document_ref(path);
match document {
Some(DocumentQuery::Text { document, .. }) => {
Notebook::from_source_code(document.contents())
}
Some(DocumentQuery::Notebook { notebook, .. }) => Ok(notebook.make_ruff_notebook()),
Some(Document::Text(document)) => Notebook::from_source_code(document.contents()),
Some(Document::Notebook(notebook)) => Ok(notebook.make_ruff_notebook()),
None => self.native_system.read_to_notebook(path),
}
}
@ -213,7 +175,7 @@ impl System for LSPSystem {
.system_virtual_path_to_document_ref(path)
.ok_or_else(|| virtual_path_not_found(path))?;
if let DocumentQuery::Text { document, .. } = &document {
if let Document::Text(document) = &document {
Ok(document.contents().to_string())
} else {
Err(not_a_text_document(path))
@ -229,8 +191,8 @@ impl System for LSPSystem {
.ok_or_else(|| virtual_path_not_found(path))?;
match document {
DocumentQuery::Text { document, .. } => Notebook::from_source_code(document.contents()),
DocumentQuery::Notebook { notebook, .. } => Ok(notebook.make_ruff_notebook()),
Document::Text(document) => Notebook::from_source_code(document.contents()),
Document::Notebook(notebook) => Ok(notebook.make_ruff_notebook()),
}
}
@ -307,7 +269,7 @@ 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: &DocumentQuery) -> FileRevision {
fn document_revision(document: &Document) -> 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)]