Add Diagnostics to make the cell diagnostics explicit

Author: dhruvmanila

crates/ty_server/src/document.rs

@@ -55,15 +55,6 @@ impl DocumentKey {
             | DocumentKey::Text(url) => url,
         }
     }
-
-    /// Converts the key back into its original URL.
-    pub(crate) fn into_url(self) -> Url {
-        match self {
-            DocumentKey::NotebookCell(url)
-            | DocumentKey::Notebook(url)
-            | DocumentKey::Text(url) => url,
-        }
-    }
 }
 
 impl std::fmt::Display for DocumentKey {

crates/ty_server/src/document/notebook.rs

@@ -211,7 +211,7 @@ impl NotebookDocument {
     }
 
     /// Returns a list of cell URIs in the order they appear in the array.
-    pub(crate) fn urls(&self) -> impl Iterator<Item = &lsp_types::Url> {
+    pub(crate) fn cell_urls(&self) -> impl Iterator<Item = &lsp_types::Url> {
         self.cells.iter().map(|cell| &cell.url)
     }
 

crates/ty_server/src/server/api/diagnostics.rs

@@ -18,8 +18,29 @@ use crate::{DocumentSnapshot, PositionEncoding};
 
 use super::LSPResult;
 
-/// A series of diagnostics across a single text document or an arbitrary number of notebook cells.
-pub(super) type DiagnosticsMap = FxHashMap<Url, Vec<Diagnostic>>;
+/// Represents the diagnostics for a text document or a notebook document.
+pub(super) enum Diagnostics {
+    TextDocument(Vec<Diagnostic>),
+
+    /// A map of cell URLs to the diagnostics for that cell.
+    NotebookDocument(FxHashMap<Url, Vec<Diagnostic>>),
+}
+
+impl Diagnostics {
+    /// Returns the diagnostics for a text document.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the diagnostics are for a notebook document.
+    pub(super) fn expect_text_document(self) -> Vec<Diagnostic> {
+        match self {
+            Diagnostics::TextDocument(diagnostics) => diagnostics,
+            Diagnostics::NotebookDocument(_) => {
+                panic!("Expected a text document diagnostics, but got notebook diagnostics")
+            }
+        }
+    }
+}
 
 /// Clears the diagnostics for the document at `uri`.
 ///
@@ -44,14 +65,29 @@ pub(super) fn publish_diagnostics_for_document(
     snapshot: &DocumentSnapshot,
     notifier: &Notifier,
 ) -> Result<()> {
-    for (uri, diagnostics) in compute_diagnostics(db, snapshot) {
+    let Some(diagnostics) = compute_diagnostics(db, snapshot) else {
+        return Ok(());
+    };
+
+    let publish_diagnostics = |uri: Url, diagnostics: Vec<Diagnostic>| {
         notifier
             .notify::<PublishDiagnostics>(PublishDiagnosticsParams {
                 uri,
                 diagnostics,
                 version: Some(snapshot.query().version()),
             })
-            .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+            .with_failure_code(lsp_server::ErrorCode::InternalError)
+    };
+
+    match diagnostics {
+        Diagnostics::TextDocument(diagnostics) => {
+            publish_diagnostics(snapshot.query().file_url().clone(), diagnostics)?;
+        }
+        Diagnostics::NotebookDocument(map) => {
+            for (cell_url, diagnostics) in map {
+                publish_diagnostics(cell_url, diagnostics)?;
+            }
+        }
     }
 
     Ok(())
@@ -60,66 +96,60 @@ pub(super) fn publish_diagnostics_for_document(
 pub(super) fn compute_diagnostics(
     db: &ProjectDatabase,
     snapshot: &DocumentSnapshot,
-) -> DiagnosticsMap {
+) -> Option<Diagnostics> {
     let Some(file) = snapshot.file(db) else {
         tracing::info!(
             "No file found for snapshot for `{}`",
             snapshot.query().file_url()
         );
-        return DiagnosticsMap::default();
+        return None;
     };
 
     let diagnostics = match db.check_file(file) {
         Ok(diagnostics) => diagnostics,
         Err(cancelled) => {
             tracing::info!("Diagnostics computation {cancelled}");
-            return DiagnosticsMap::default();
+            return None;
         }
     };
 
-    let mut diagnostics_map = DiagnosticsMap::default();
     let query = snapshot.query();
-    // let source_kind = query.make_source_kind();
 
-    // Populates all relevant URLs with an empty diagnostic list.
-    // This ensures that documents without diagnostics still get updated.
     if let Some(notebook) = query.as_notebook() {
-        for url in notebook.urls() {
-            diagnostics_map.entry(url.clone()).or_default();
-        }
-    } else {
-        diagnostics_map
-            .entry(query.make_key().into_url())
-            .or_default();
-    }
+        let mut cell_diagnostics: FxHashMap<Url, Vec<Diagnostic>> = FxHashMap::default();
 
-    let lsp_diagnostics = diagnostics.as_slice().iter().map(|diagnostic| {
-        (
-            // TODO: Use the cell index instead using `source_kind`
-            usize::default(),
-            to_lsp_diagnostic(db, diagnostic, snapshot.encoding()),
-        )
-    });
+        // Populates all relevant URLs with an empty diagnostic list. This ensures that documents
+        // without diagnostics still get updated.
+        for cell_url in notebook.cell_urls() {
+            cell_diagnostics.entry(cell_url.clone()).or_default();
+        }
 
-    if let Some(notebook) = query.as_notebook() {
-        for (index, diagnostic) in lsp_diagnostics {
-            let Some(uri) = notebook.cell_uri_by_index(index) else {
-                tracing::warn!("Unable to find notebook cell at index {index}");
+        for (cell_index, diagnostic) in diagnostics.iter().map(|diagnostic| {
+            (
+                // TODO: Use the cell index instead using `SourceKind`
+                usize::default(),
+                to_lsp_diagnostic(db, diagnostic, snapshot.encoding()),
+            )
+        }) {
+            let Some(cell_uri) = notebook.cell_uri_by_index(cell_index) else {
+                tracing::warn!("Unable to find notebook cell at index {cell_index}");
                 continue;
             };
-            diagnostics_map
-                .entry(uri.clone())
+            cell_diagnostics
+                .entry(cell_uri.clone())
                 .or_default()
                 .push(diagnostic);
         }
+
+        Some(Diagnostics::NotebookDocument(cell_diagnostics))
     } else {
-        diagnostics_map
-            .entry(query.make_key().into_url())
-            .or_default()
-            .extend(lsp_diagnostics.map(|(_, diagnostic)| diagnostic));
+        Some(Diagnostics::TextDocument(
+            diagnostics
+                .iter()
+                .map(|diagnostic| to_lsp_diagnostic(db, diagnostic, snapshot.encoding()))
+                .collect(),
+        ))
     }
-
-    diagnostics_map
 }
 
 /// Converts the tool specific [`Diagnostic`][ruff_db::diagnostic::Diagnostic] to an LSP

crates/ty_server/src/server/api/requests/diagnostic.rs

@@ -6,7 +6,7 @@ use lsp_types::{
     FullDocumentDiagnosticReport, RelatedFullDocumentDiagnosticReport,
 };
 
-use crate::server::api::diagnostics::compute_diagnostics;
+use crate::server::api::diagnostics::{Diagnostics, compute_diagnostics};
 use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
 use crate::server::{Result, client::Notifier};
 use crate::session::DocumentSnapshot;
@@ -34,13 +34,10 @@ impl BackgroundDocumentRequestHandler for DocumentDiagnosticRequestHandler {
                 related_documents: None,
                 full_document_diagnostic_report: FullDocumentDiagnosticReport {
                     result_id: None,
-                    // Pull diagnostic requests are only called for text documents, not for
+                    // SAFETY: Pull diagnostic requests are only called for text documents, not for
                     // notebook documents.
                     items: compute_diagnostics(db, &snapshot)
-                        .into_iter()
-                        .next()
-                        .map(|(_, diagnostics)| diagnostics)
-                        .unwrap_or_default(),
+                        .map_or_else(Vec::new, Diagnostics::expect_text_document),
                 },
             }),
         ))

crates/ty_server/src/session/index.rs

@@ -135,7 +135,7 @@ impl Index {
     }
 
     pub(super) fn open_notebook_document(&mut self, notebook_url: Url, document: NotebookDocument) {
-        for cell_url in document.urls() {
+        for cell_url in document.cell_urls() {
             self.notebook_cells
                 .insert(cell_url.clone(), notebook_url.clone());
         }
@@ -265,6 +265,7 @@ 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()),