diff --git a/crates/ty_server/src/document.rs b/crates/ty_server/src/document.rs index dade6a7fd1..fff51d2f49 100644 --- a/crates/ty_server/src/document.rs +++ b/crates/ty_server/src/document.rs @@ -7,6 +7,8 @@ mod text_document; pub(crate) use location::ToLink; use lsp_types::{PositionEncodingKind, Url}; + +use crate::system::AnySystemPath; pub use notebook::NotebookDocument; pub(crate) use range::{FileRangeExt, PositionExt, RangeExt, TextSizeExt, ToRangeExt}; pub(crate) use text_document::DocumentVersion; @@ -40,19 +42,38 @@ impl From 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 enum DocumentKey { - Notebook(Url), - NotebookCell(Url), - Text(Url), +pub(crate) enum DocumentKey { + Notebook(AnySystemPath), + NotebookCell { + cell_url: Url, + notebook_path: AnySystemPath, + }, + Text(AnySystemPath), } impl DocumentKey { - /// Returns the URL associated with the key. - pub(crate) fn url(&self) -> &Url { + /// Returns the file path associated with the key. + pub(crate) fn path(&self) -> &AnySystemPath { match self { - DocumentKey::NotebookCell(url) - | DocumentKey::Notebook(url) - | DocumentKey::Text(url) => url, + DocumentKey::Notebook(path) | DocumentKey::Text(path) => path, + DocumentKey::NotebookCell { notebook_path, .. } => notebook_path, + } + } + + 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), + } + } + + /// 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 { + match self { + DocumentKey::NotebookCell { cell_url, .. } => Some(cell_url.clone()), + DocumentKey::Notebook(path) | DocumentKey::Text(path) => path.to_url(), } } } @@ -60,7 +81,11 @@ impl DocumentKey { impl std::fmt::Display for DocumentKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - Self::NotebookCell(url) | Self::Notebook(url) | Self::Text(url) => url.fmt(f), + 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), + }, } } } diff --git a/crates/ty_server/src/lib.rs b/crates/ty_server/src/lib.rs index 5589c60464..cd7d71a73c 100644 --- a/crates/ty_server/src/lib.rs +++ b/crates/ty_server/src/lib.rs @@ -1,6 +1,6 @@ use crate::server::{ConnectionInitializer, Server}; use anyhow::Context; -pub use document::{DocumentKey, NotebookDocument, PositionEncoding, TextDocument}; +pub use document::{NotebookDocument, PositionEncoding, TextDocument}; pub use session::{ClientSettings, DocumentQuery, DocumentSnapshot, Session}; use std::num::NonZeroUsize; diff --git a/crates/ty_server/src/server/api.rs b/crates/ty_server/src/server/api.rs index db0c1bb1a7..f937adb2e4 100644 --- a/crates/ty_server/src/server/api.rs +++ b/crates/ty_server/src/server/api.rs @@ -1,6 +1,6 @@ use crate::server::schedule::Task; use crate::session::Session; -use crate::system::{AnySystemPath, url_to_any_system_path}; +use crate::system::AnySystemPath; use anyhow::anyhow; use lsp_server as server; use lsp_server::RequestId; @@ -154,7 +154,7 @@ where let url = R::document_url(¶ms).into_owned(); - let Ok(path) = url_to_any_system_path(&url) else { + let Ok(path) = AnySystemPath::try_from_url(&url) else { tracing::warn!("Ignoring request for invalid `{url}`"); return Box::new(|_| {}); }; diff --git a/crates/ty_server/src/server/api/diagnostics.rs b/crates/ty_server/src/server/api/diagnostics.rs index 57ffbd6650..812a11f9a8 100644 --- a/crates/ty_server/src/server/api/diagnostics.rs +++ b/crates/ty_server/src/server/api/diagnostics.rs @@ -12,10 +12,9 @@ use ruff_db::source::{line_index, source_text}; use ty_project::{Db, ProjectDatabase}; use super::LSPResult; -use crate::document::{FileRangeExt, ToRangeExt}; +use crate::document::{DocumentKey, FileRangeExt, ToRangeExt}; use crate::server::Result; use crate::session::client::Client; -use crate::system::url_to_any_system_path; use crate::{DocumentSnapshot, PositionEncoding, Session}; /// Represents the diagnostics for a text document or a notebook document. @@ -42,13 +41,20 @@ impl Diagnostics { } } -/// Clears the diagnostics for the document at `uri`. +/// Clears the diagnostics for the document identified by `key`. /// /// This is done by notifying the client with an empty list of diagnostics for the document. -pub(super) fn clear_diagnostics(uri: &Url, client: &Client) -> Result<()> { +/// 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(key: &DocumentKey, client: &Client) -> Result<()> { + let Some(uri) = key.to_url() else { + // If we can't convert to URL, we can't clear diagnostics + return Ok(()); + }; + client .send_notification::(PublishDiagnosticsParams { - uri: uri.clone(), + uri, diagnostics: vec![], version: None, }) @@ -62,21 +68,27 @@ pub(super) fn clear_diagnostics(uri: &Url, client: &Client) -> Result<()> { /// 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, url: Url, client: &Client) -> Result<()> { +pub(super) fn publish_diagnostics( + session: &Session, + key: &DocumentKey, + client: &Client, +) -> Result<()> { if session.client_capabilities().pull_diagnostics { return Ok(()); } - let Ok(path) = url_to_any_system_path(&url) else { + let Some(url) = key.to_url() else { return Ok(()); }; + let path = key.path(); + let snapshot = session .take_snapshot(url.clone()) .ok_or_else(|| anyhow::anyhow!("Unable to take snapshot for document with URL {url}")) .with_failure_code(lsp_server::ErrorCode::InternalError)?; - let db = session.project_db_or_default(&path); + let db = session.project_db_or_default(path); let Some(diagnostics) = compute_diagnostics(db, &snapshot) else { return Ok(()); diff --git a/crates/ty_server/src/server/api/notifications/did_change.rs b/crates/ty_server/src/server/api/notifications/did_change.rs index 43b8f88c7d..b7df7961c4 100644 --- a/crates/ty_server/src/server/api/notifications/did_change.rs +++ b/crates/ty_server/src/server/api/notifications/did_change.rs @@ -8,7 +8,7 @@ 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, url_to_any_system_path}; +use crate::system::AnySystemPath; use ty_project::watch::ChangeEvent; pub(crate) struct DidChangeTextDocumentHandler; @@ -28,30 +28,32 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler { content_changes, } = params; - let Ok(path) = url_to_any_system_path(&uri) else { + let Ok(key) = session.key_from_url(uri.clone()) else { + tracing::debug!("Failed to create document key from URI: {}", uri); return Ok(()); }; - let key = session.key_from_url(uri.clone()); - session .update_text_document(&key, content_changes, version) .with_failure_code(ErrorCode::InternalError)?; - match path { + match key.path() { AnySystemPath::System(path) => { let db = match session.project_db_for_path_mut(path.as_std_path()) { Some(db) => db, None => session.default_project_db_mut(), }; - db.apply_changes(vec![ChangeEvent::file_content_changed(path)], None); + db.apply_changes(vec![ChangeEvent::file_content_changed(path.clone())], None); } AnySystemPath::SystemVirtual(virtual_path) => { let db = session.default_project_db_mut(); - db.apply_changes(vec![ChangeEvent::ChangedVirtual(virtual_path)], None); + db.apply_changes( + vec![ChangeEvent::ChangedVirtual(virtual_path.clone())], + None, + ); } } - publish_diagnostics(session, uri, client) + publish_diagnostics(session, &key, client) } } diff --git a/crates/ty_server/src/server/api/notifications/did_change_watched_files.rs b/crates/ty_server/src/server/api/notifications/did_change_watched_files.rs index 2d27495cfe..98c67829b1 100644 --- a/crates/ty_server/src/server/api/notifications/did_change_watched_files.rs +++ b/crates/ty_server/src/server/api/notifications/did_change_watched_files.rs @@ -4,7 +4,7 @@ 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, url_to_any_system_path}; +use crate::system::AnySystemPath; use lsp_types as types; use lsp_types::{FileChangeType, notification as notif}; use rustc_hash::FxHashMap; @@ -26,7 +26,7 @@ impl SyncNotificationHandler for DidChangeWatchedFiles { let mut events_by_db: FxHashMap<_, Vec> = FxHashMap::default(); for change in params.changes { - let path = match url_to_any_system_path(&change.uri) { + let path = match AnySystemPath::try_from_url(&change.uri) { Ok(path) => path, Err(err) => { tracing::warn!( @@ -111,8 +111,8 @@ impl SyncNotificationHandler for DidChangeWatchedFiles { ) .with_failure_code(lsp_server::ErrorCode::InternalError)?; } else { - for url in session.text_document_urls() { - publish_diagnostics(session, url.clone(), client)?; + for key in session.text_document_keys() { + publish_diagnostics(session, &key, client)?; } } diff --git a/crates/ty_server/src/server/api/notifications/did_close.rs b/crates/ty_server/src/server/api/notifications/did_close.rs index 649546595e..b7b06cdc82 100644 --- a/crates/ty_server/src/server/api/notifications/did_close.rs +++ b/crates/ty_server/src/server/api/notifications/did_close.rs @@ -4,7 +4,7 @@ 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, url_to_any_system_path}; +use crate::system::AnySystemPath; use lsp_server::ErrorCode; use lsp_types::DidCloseTextDocumentParams; use lsp_types::notification::DidCloseTextDocument; @@ -22,22 +22,25 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler { client: &Client, params: DidCloseTextDocumentParams, ) -> Result<()> { - let Ok(path) = url_to_any_system_path(¶ms.text_document.uri) else { + let Ok(key) = session.key_from_url(params.text_document.uri.clone()) else { + tracing::debug!( + "Failed to create document key from URI: {}", + params.text_document.uri + ); return Ok(()); }; - - let key = session.key_from_url(params.text_document.uri); session .close_document(&key) .with_failure_code(ErrorCode::InternalError)?; - if let AnySystemPath::SystemVirtual(virtual_path) = path { + if let AnySystemPath::SystemVirtual(virtual_path) = key.path() { let db = session.default_project_db_mut(); - db.apply_changes(vec![ChangeEvent::DeletedVirtual(virtual_path)], None); + db.apply_changes( + vec![ChangeEvent::DeletedVirtual(virtual_path.clone())], + None, + ); } - clear_diagnostics(key.url(), client)?; - - Ok(()) + clear_diagnostics(&key, client) } } diff --git a/crates/ty_server/src/server/api/notifications/did_close_notebook.rs b/crates/ty_server/src/server/api/notifications/did_close_notebook.rs index 627e5415e3..6916e56c32 100644 --- a/crates/ty_server/src/server/api/notifications/did_close_notebook.rs +++ b/crates/ty_server/src/server/api/notifications/did_close_notebook.rs @@ -6,7 +6,7 @@ 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, url_to_any_system_path}; +use crate::system::AnySystemPath; use ty_project::watch::ChangeEvent; pub(crate) struct DidCloseNotebookHandler; @@ -21,18 +21,23 @@ impl SyncNotificationHandler for DidCloseNotebookHandler { _client: &Client, params: DidCloseNotebookDocumentParams, ) -> Result<()> { - let Ok(path) = url_to_any_system_path(¶ms.notebook_document.uri) else { + let Ok(key) = session.key_from_url(params.notebook_document.uri.clone()) else { + tracing::debug!( + "Failed to create document key from URI: {}", + params.notebook_document.uri + ); return Ok(()); }; - - let key = session.key_from_url(params.notebook_document.uri); session .close_document(&key) .with_failure_code(lsp_server::ErrorCode::InternalError)?; - if let AnySystemPath::SystemVirtual(virtual_path) = path { + if let AnySystemPath::SystemVirtual(virtual_path) = key.path() { let db = session.default_project_db_mut(); - db.apply_changes(vec![ChangeEvent::DeletedVirtual(virtual_path)], None); + db.apply_changes( + vec![ChangeEvent::DeletedVirtual(virtual_path.clone())], + None, + ); } Ok(()) diff --git a/crates/ty_server/src/server/api/notifications/did_open.rs b/crates/ty_server/src/server/api/notifications/did_open.rs index 9f2dea3691..1f61d07f7c 100644 --- a/crates/ty_server/src/server/api/notifications/did_open.rs +++ b/crates/ty_server/src/server/api/notifications/did_open.rs @@ -7,7 +7,7 @@ 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, url_to_any_system_path}; +use crate::system::AnySystemPath; use ruff_db::Db; use ty_project::watch::ChangeEvent; @@ -31,20 +31,21 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler { }, }: DidOpenTextDocumentParams, ) -> Result<()> { - let Ok(path) = url_to_any_system_path(&uri) else { + let Ok(key) = session.key_from_url(uri.clone()) else { + tracing::debug!("Failed to create document key from URI: {}", uri); return Ok(()); }; let document = TextDocument::new(text, version).with_language_id(&language_id); - session.open_text_document(uri.clone(), document); + session.open_text_document(key.path(), document); - match &path { - AnySystemPath::System(path) => { - let db = match session.project_db_for_path_mut(path.as_std_path()) { + match key.path() { + AnySystemPath::System(system_path) => { + let db = match session.project_db_for_path_mut(system_path.as_std_path()) { Some(db) => db, None => session.default_project_db_mut(), }; - db.apply_changes(vec![ChangeEvent::Opened(path.clone())], None); + db.apply_changes(vec![ChangeEvent::Opened(system_path.clone())], None); } AnySystemPath::SystemVirtual(virtual_path) => { let db = session.default_project_db_mut(); @@ -52,6 +53,6 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler { } } - publish_diagnostics(session, uri, client) + publish_diagnostics(session, &key, client) } } diff --git a/crates/ty_server/src/server/api/notifications/did_open_notebook.rs b/crates/ty_server/src/server/api/notifications/did_open_notebook.rs index 9a8d24bb9c..ccae65a15d 100644 --- a/crates/ty_server/src/server/api/notifications/did_open_notebook.rs +++ b/crates/ty_server/src/server/api/notifications/did_open_notebook.rs @@ -11,7 +11,7 @@ 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, url_to_any_system_path}; +use crate::system::AnySystemPath; pub(crate) struct DidOpenNotebookHandler; @@ -25,7 +25,7 @@ impl SyncNotificationHandler for DidOpenNotebookHandler { _client: &Client, params: DidOpenNotebookDocumentParams, ) -> Result<()> { - let Ok(path) = url_to_any_system_path(¶ms.notebook_document.uri) else { + let Ok(path) = AnySystemPath::try_from_url(¶ms.notebook_document.uri) else { return Ok(()); }; @@ -36,19 +36,19 @@ impl SyncNotificationHandler for DidOpenNotebookHandler { params.cell_text_documents, ) .with_failure_code(ErrorCode::InternalError)?; - session.open_notebook_document(params.notebook_document.uri, notebook); + session.open_notebook_document(&path, notebook); - match path { - AnySystemPath::System(path) => { - let db = match session.project_db_for_path_mut(path.as_std_path()) { + match &path { + AnySystemPath::System(system_path) => { + let db = match session.project_db_for_path_mut(system_path.as_std_path()) { Some(db) => db, None => session.default_project_db_mut(), }; - db.apply_changes(vec![ChangeEvent::Opened(path)], None); + db.apply_changes(vec![ChangeEvent::Opened(system_path.clone())], None); } AnySystemPath::SystemVirtual(virtual_path) => { let db = session.default_project_db_mut(); - db.files().virtual_file(db, &virtual_path); + db.files().virtual_file(db, virtual_path); } } diff --git a/crates/ty_server/src/session.rs b/crates/ty_server/src/session.rs index cced715ad5..f488ddecc9 100644 --- a/crates/ty_server/src/session.rs +++ b/crates/ty_server/src/session.rs @@ -19,7 +19,7 @@ pub use self::settings::ClientSettings; pub(crate) use self::settings::Experimental; use crate::document::{DocumentKey, DocumentVersion, NotebookDocument}; use crate::session::request_queue::RequestQueue; -use crate::system::{AnySystemPath, LSPSystem, url_to_any_system_path}; +use crate::system::{AnySystemPath, LSPSystem}; use crate::{PositionEncoding, TextDocument}; mod capabilities; @@ -158,35 +158,43 @@ impl Session { .unwrap() } - pub fn key_from_url(&self, url: Url) -> DocumentKey { + pub(crate) fn key_from_url(&self, url: Url) -> crate::Result { self.index().key_from_url(url) } /// Creates a document snapshot with the URL referencing the document to snapshot. + /// + /// Returns `None` if the url can't be converted to a document key or if the document isn't open. pub fn take_snapshot(&self, url: Url) -> Option { - let key = self.key_from_url(url); + let key = self.key_from_url(url).ok()?; Some(DocumentSnapshot { resolved_client_capabilities: self.resolved_client_capabilities.clone(), - document_ref: self.index().make_document_ref(key)?, + document_ref: self.index().make_document_ref(&key)?, position_encoding: self.position_encoding, }) } - /// Iterates over the LSP URLs for all open text documents. These URLs are valid file paths. - pub(super) fn text_document_urls(&self) -> impl Iterator + '_ { - self.index().text_document_urls() + /// Iterates over the document keys for all open text documents. + pub(super) fn text_document_keys(&self) -> impl Iterator + '_ { + self.index() + .text_document_paths() + .map(|path| DocumentKey::Text(path.clone())) } - /// Registers a notebook document at the provided `url`. + /// Registers a notebook document at the provided `path`. /// If a document is already open here, it will be overwritten. - pub fn open_notebook_document(&mut self, url: Url, document: NotebookDocument) { - self.index_mut().open_notebook_document(url, document); + pub(crate) fn open_notebook_document( + &mut self, + path: &AnySystemPath, + document: NotebookDocument, + ) { + self.index_mut().open_notebook_document(path, document); } - /// Registers a text document at the provided `url`. + /// 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, url: Url, document: TextDocument) { - self.index_mut().open_text_document(url, document); + pub(crate) fn open_text_document(&mut self, path: &AnySystemPath, document: TextDocument) { + self.index_mut().open_text_document(path, document); } /// Updates a text document at the associated `key`. @@ -314,7 +322,7 @@ impl DocumentSnapshot { } pub(crate) fn file(&self, db: &dyn Db) -> Option { - match url_to_any_system_path(self.document_ref.file_url()).ok()? { + match AnySystemPath::try_from_url(self.document_ref.file_url()).ok()? { AnySystemPath::System(path) => system_path_to_file(db, path).ok(), AnySystemPath::SystemVirtual(virtual_path) => db .files() diff --git a/crates/ty_server/src/session/index.rs b/crates/ty_server/src/session/index.rs index da0b7570ef..3a15b1c9a3 100644 --- a/crates/ty_server/src/session/index.rs +++ b/crates/ty_server/src/session/index.rs @@ -1,4 +1,3 @@ -use std::path::Path; use std::sync::Arc; use lsp_types::Url; @@ -7,6 +6,7 @@ use rustc_hash::FxHashMap; use crate::{ PositionEncoding, TextDocument, document::{DocumentKey, DocumentVersion, NotebookDocument}, + system::AnySystemPath, }; use super::ClientSettings; @@ -14,11 +14,11 @@ use super::ClientSettings; /// Stores and tracks all open documents in a session, along with their associated settings. #[derive(Default, Debug)] pub(crate) struct Index { - /// Maps all document file URLs to the associated document controller - documents: FxHashMap, + /// Maps all document file paths to the associated document controller + documents: FxHashMap, - /// Maps opaque cell URLs to a notebook URL (document) - notebook_cells: FxHashMap, + /// Maps opaque cell URLs to a notebook path (document) + notebook_cells: FxHashMap, /// Global settings provided by the client. #[expect(dead_code)] @@ -34,18 +34,18 @@ impl Index { } } - pub(super) fn text_document_urls(&self) -> impl Iterator + '_ { + pub(super) fn text_document_paths(&self) -> impl Iterator + '_ { self.documents .iter() - .filter_map(|(url, doc)| doc.as_text().and(Some(url))) + .filter_map(|(path, doc)| doc.as_text().and(Some(path))) } #[expect(dead_code)] - pub(super) fn notebook_document_urls(&self) -> impl Iterator + '_ { + pub(super) fn notebook_document_paths(&self) -> impl Iterator + '_ { self.documents .iter() .filter(|(_, doc)| doc.as_notebook().is_some()) - .map(|(url, _)| url) + .map(|(path, _)| path) } pub(super) fn update_text_document( @@ -57,7 +57,7 @@ impl Index { ) -> crate::Result<()> { let controller = self.document_controller_for_key(key)?; let Some(document) = controller.as_text_mut() else { - anyhow::bail!("Text document URI does not point to a text document"); + anyhow::bail!("Text document path does not point to a text document"); }; if content_changes.is_empty() { @@ -70,16 +70,24 @@ impl Index { Ok(()) } - pub(crate) fn key_from_url(&self, url: Url) -> DocumentKey { - if self.notebook_cells.contains_key(&url) { - DocumentKey::NotebookCell(url) - } else if Path::new(url.path()) - .extension() - .is_some_and(|ext| ext.eq_ignore_ascii_case("ipynb")) - { - DocumentKey::Notebook(url) + pub(crate) fn key_from_url(&self, url: Url) -> crate::Result { + if let Some(notebook_path) = self.notebook_cells.get(&url) { + Ok(DocumentKey::NotebookCell { + cell_url: url, + notebook_path: notebook_path.clone(), + }) } else { - DocumentKey::Text(url) + let path = AnySystemPath::try_from_url(&url) + .map_err(|()| anyhow::anyhow!("Failed to convert URL to system path: {}", url))?; + + if path + .extension() + .is_some_and(|ext| ext.eq_ignore_ascii_case("ipynb")) + { + Ok(DocumentKey::Notebook(path)) + } else { + Ok(DocumentKey::Text(path)) + } } } @@ -98,48 +106,55 @@ impl Index { .. }) = cells.as_ref().and_then(|cells| cells.structure.as_ref()) { - let Some(path) = self.url_for_key(key).cloned() else { - anyhow::bail!("Tried to open unavailable document `{key}`"); - }; + let notebook_path = key.path().clone(); for opened_cell in did_open { self.notebook_cells - .insert(opened_cell.uri.clone(), path.clone()); + .insert(opened_cell.uri.clone(), notebook_path.clone()); } // 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 { - anyhow::bail!("Notebook document URI does not point to a notebook document"); + anyhow::bail!("Notebook document path does not point to a notebook document"); }; notebook.update(cells, metadata, new_version, encoding)?; Ok(()) } - pub(crate) fn make_document_ref(&self, key: DocumentKey) -> Option { - let url = self.url_for_key(&key)?.clone(); - let controller = self.documents.get(&url)?; - let cell_url = match key { - DocumentKey::NotebookCell(cell_url) => Some(cell_url), - _ => None, + pub(crate) fn make_document_ref(&self, key: &DocumentKey) -> Option { + let path = key.path(); + let controller = self.documents.get(path)?; + let (cell_url, file_url) = match &key { + DocumentKey::NotebookCell { + cell_url, + notebook_path, + } => (Some(cell_url.clone()), notebook_path.to_url()?), + DocumentKey::Notebook(path) | DocumentKey::Text(path) => (None, path.to_url()?), }; - Some(controller.make_ref(cell_url, url)) + Some(controller.make_ref(cell_url, file_url)) } - pub(super) fn open_text_document(&mut self, url: Url, document: TextDocument) { + pub(super) fn open_text_document(&mut self, path: &AnySystemPath, document: TextDocument) { self.documents - .insert(url, DocumentController::new_text(document)); + .insert(path.clone(), DocumentController::new_text(document)); } - pub(super) fn open_notebook_document(&mut self, notebook_url: Url, document: NotebookDocument) { + pub(super) fn open_notebook_document( + &mut self, + notebook_path: &AnySystemPath, + document: NotebookDocument, + ) { for cell_url in document.cell_urls() { self.notebook_cells - .insert(cell_url.clone(), notebook_url.clone()); + .insert(cell_url.clone(), notebook_path.clone()); } - self.documents - .insert(notebook_url, DocumentController::new_notebook(document)); + self.documents.insert( + notebook_path.clone(), + DocumentController::new_notebook(document), + ); } pub(super) fn close_document(&mut self, key: &DocumentKey) -> crate::Result<()> { @@ -148,18 +163,16 @@ 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(uri) = key { - if self.notebook_cells.remove(uri).is_none() { - tracing::warn!("Tried to remove a notebook cell that does not exist: {uri}",); + 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(()); } - let Some(url) = self.url_for_key(key).cloned() else { - anyhow::bail!("Tried to close unavailable document `{key}`"); - }; + let path = key.path(); - let Some(_) = self.documents.remove(&url) else { - anyhow::bail!("tried to close document that didn't exist at {}", url) + let Some(_) = self.documents.remove(path) else { + anyhow::bail!("tried to close document that didn't exist at {}", key) }; Ok(()) } @@ -168,21 +181,12 @@ impl Index { &mut self, key: &DocumentKey, ) -> crate::Result<&mut DocumentController> { - let Some(url) = self.url_for_key(key).cloned() else { - anyhow::bail!("Tried to open unavailable document `{key}`"); - }; - let Some(controller) = self.documents.get_mut(&url) else { - anyhow::bail!("Document controller not available at `{}`", url); + let path = key.path(); + let Some(controller) = self.documents.get_mut(path) else { + anyhow::bail!("Document controller not available at `{}`", key); }; Ok(controller) } - - fn url_for_key<'a>(&'a self, key: &'a DocumentKey) -> Option<&'a Url> { - match key { - DocumentKey::Notebook(path) | DocumentKey::Text(path) => Some(path), - DocumentKey::NotebookCell(uri) => self.notebook_cells.get(uri), - } - } } /// A mutable handler to an underlying document. @@ -263,19 +267,6 @@ pub enum DocumentQuery { } impl DocumentQuery { - /// Retrieve the original key that describes this document query. - #[expect(dead_code)] - pub(crate) fn make_key(&self) -> DocumentKey { - match self { - Self::Text { file_url, .. } => DocumentKey::Text(file_url.clone()), - Self::Notebook { - cell_url: Some(cell_uri), - .. - } => DocumentKey::NotebookCell(cell_uri.clone()), - Self::Notebook { file_url, .. } => DocumentKey::Notebook(file_url.clone()), - } - } - /// Attempts to access the underlying notebook document that this query is selecting. pub fn as_notebook(&self) -> Option<&NotebookDocument> { match self { diff --git a/crates/ty_server/src/system.rs b/crates/ty_server/src/system.rs index 4a38ed9046..e946029ff3 100644 --- a/crates/ty_server/src/system.rs +++ b/crates/ty_server/src/system.rs @@ -14,28 +14,9 @@ use ruff_notebook::{Notebook, NotebookError}; use ty_python_semantic::Db; use crate::DocumentQuery; +use crate::document::DocumentKey; use crate::session::index::Index; -/// 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 url_to_any_system_path(url: &Url) -> std::result::Result { - 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) fn file_to_url(db: &dyn Db, file: File) -> Option { match file.path(db) { FilePath::System(system) => Url::from_file_path(system.as_std_path()).ok(), @@ -47,19 +28,57 @@ pub(crate) fn file_to_url(db: &dyn Db, file: File) -> Option { } /// Represents either a [`SystemPath`] or a [`SystemVirtualPath`]. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Hash, PartialEq, Eq)] pub(crate) enum AnySystemPath { System(SystemPathBuf), SystemVirtual(SystemVirtualPathBuf), } 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 { + 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), AnySystemPath::SystemVirtual(_) => None, } } + + /// Returns the extension of the path, if any. + pub(crate) fn extension(&self) -> Option<&str> { + 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 { + 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(), + } + } } #[derive(Debug)] @@ -107,39 +126,29 @@ impl LSPSystem { self.index.as_ref().unwrap() } - fn make_document_ref(&self, url: Url) -> Option { + fn make_document_ref(&self, path: AnySystemPath) -> Option { let index = self.index(); - let key = index.key_from_url(url); - index.make_document_ref(key) + let key = DocumentKey::from_path(path); + index.make_document_ref(&key) } - fn system_path_to_document_ref(&self, path: &SystemPath) -> Result> { - let url = Url::from_file_path(path.as_std_path()).map_err(|()| { - std::io::Error::new( - std::io::ErrorKind::InvalidInput, - format!("Failed to convert system path to URL: {path:?}"), - ) - })?; - Ok(self.make_document_ref(url)) + fn system_path_to_document_ref(&self, path: &SystemPath) -> Option { + 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, - ) -> Result> { - let url = Url::parse(path.as_str()).map_err(|_| { - std::io::Error::new( - std::io::ErrorKind::InvalidInput, - format!("Failed to convert virtual path to URL: {path:?}"), - ) - })?; - Ok(self.make_document_ref(url)) + ) -> Option { + let any_path = AnySystemPath::SystemVirtual(path.to_path_buf()); + self.make_document_ref(any_path) } } impl System for LSPSystem { fn path_metadata(&self, path: &SystemPath) -> Result { - let document = self.system_path_to_document_ref(path)?; + let document = self.system_path_to_document_ref(path); if let Some(document) = document { Ok(Metadata::new( @@ -161,7 +170,7 @@ impl System for LSPSystem { } fn read_to_string(&self, path: &SystemPath) -> Result { - let document = self.system_path_to_document_ref(path)?; + let document = self.system_path_to_document_ref(path); match document { Some(DocumentQuery::Text { document, .. }) => Ok(document.contents().to_string()), @@ -170,7 +179,7 @@ impl System for LSPSystem { } fn read_to_notebook(&self, path: &SystemPath) -> std::result::Result { - let document = self.system_path_to_document_ref(path)?; + let document = self.system_path_to_document_ref(path); match document { Some(DocumentQuery::Text { document, .. }) => { @@ -183,7 +192,7 @@ impl System for LSPSystem { fn read_virtual_path_to_string(&self, path: &SystemVirtualPath) -> Result { let document = self - .system_virtual_path_to_document_ref(path)? + .system_virtual_path_to_document_ref(path) .ok_or_else(|| virtual_path_not_found(path))?; if let DocumentQuery::Text { document, .. } = &document { @@ -198,7 +207,7 @@ impl System for LSPSystem { path: &SystemVirtualPath, ) -> std::result::Result { let document = self - .system_virtual_path_to_document_ref(path)? + .system_virtual_path_to_document_ref(path) .ok_or_else(|| virtual_path_not_found(path))?; match document {