[ty] Support publishing diagnostics in the server

crates/ty_project/src/db/changes.rs

@@ -11,13 +11,31 @@ use ruff_db::system::SystemPath;
 use rustc_hash::FxHashSet;
 use ty_python_semantic::Program;
 
+/// Represents the result of applying changes to the project database.
+pub struct ChangeResult {
+    project_changed: bool,
+    custom_stdlib_changed: bool,
+}
+
+impl ChangeResult {
+    /// Returns `true` if the project structure has changed.
+    pub fn project_changed(&self) -> bool {
+        self.project_changed
+    }
+
+    /// Returns `true` if the custom stdlib's VERSIONS file has changed.
+    pub fn custom_stdlib_changed(&self) -> bool {
+        self.custom_stdlib_changed
+    }
+}
+
 impl ProjectDatabase {
     #[tracing::instrument(level = "debug", skip(self, changes, project_options_overrides))]
     pub fn apply_changes(
         &mut self,
         changes: Vec<ChangeEvent>,
         project_options_overrides: Option<&ProjectOptionsOverrides>,
-    ) {
+    ) -> ChangeResult {
         let mut project = self.project();
         let project_root = project.root(self).to_path_buf();
         let config_file_override =
@@ -29,10 +47,10 @@ impl ProjectDatabase {
             .custom_stdlib_search_path(self)
             .map(|path| path.join("VERSIONS"));
 
-        // Are there structural changes to the project
-        let mut project_changed = false;
-        // Changes to a custom stdlib path's VERSIONS
-        let mut custom_stdlib_change = false;
+        let mut result = ChangeResult {
+            project_changed: false,
+            custom_stdlib_changed: false,
+        };
         // Paths that were added
         let mut added_paths = FxHashSet::default();
 
@@ -52,7 +70,7 @@ impl ProjectDatabase {
             if let Some(path) = change.system_path() {
                 if let Some(config_file) = &config_file_override {
                     if config_file.as_path() == path {
-                        project_changed = true;
+                        result.project_changed = true;
 
                         continue;
                     }
@@ -63,13 +81,13 @@ impl ProjectDatabase {
                     Some(".gitignore" | ".ignore" | "ty.toml" | "pyproject.toml")
                 ) {
                     // Changes to ignore files or settings can change the project structure or add/remove files.
-                    project_changed = true;
+                    result.project_changed = true;
 
                     continue;
                 }
 
                 if Some(path) == custom_stdlib_versions_path.as_deref() {
-                    custom_stdlib_change = true;
+                    result.custom_stdlib_changed = true;
                 }
             }
 
@@ -132,7 +150,7 @@ impl ProjectDatabase {
                             .as_ref()
                             .is_some_and(|versions_path| versions_path.starts_with(&path))
                         {
-                            custom_stdlib_change = true;
+                            result.custom_stdlib_changed = true;
                         }
 
                         if project.is_path_included(self, &path) || path == project_root {
@@ -146,7 +164,7 @@ impl ProjectDatabase {
                             // We may want to make this more clever in the future, to e.g. iterate over the
                             // indexed files and remove the once that start with the same path, unless
                             // the deleted path is the project configuration.
-                            project_changed = true;
+                            result.project_changed = true;
                         }
                     }
                 }
@@ -162,7 +180,7 @@ impl ProjectDatabase {
                 }
 
                 ChangeEvent::Rescan => {
-                    project_changed = true;
+                    result.project_changed = true;
                     Files::sync_all(self);
                     sync_recursively.clear();
                     break;
@@ -185,7 +203,7 @@ impl ProjectDatabase {
             last = Some(path);
         }
 
-        if project_changed {
+        if result.project_changed {
             let new_project_metadata = match config_file_override {
                 Some(config_file) => ProjectMetadata::from_config_file(config_file, self.system()),
                 None => ProjectMetadata::discover(&project_root, self.system()),
@@ -227,8 +245,8 @@ impl ProjectDatabase {
                 }
             }
 
-            return;
-        } else if custom_stdlib_change {
+            return result;
+        } else if result.custom_stdlib_changed {
             let search_paths = project
                 .metadata(self)
                 .to_program_settings(self.system())
@@ -258,5 +276,7 @@ impl ProjectDatabase {
         // implement a `BTreeMap` or similar and only prune the diagnostics from paths that we've
         // re-scanned (or that were removed etc).
         project.replace_index_diagnostics(self, diagnostics);
+
+        result
     }
 }

crates/ty_server/src/document/notebook.rs

@@ -199,7 +199,6 @@ impl NotebookDocument {
     }
 
     /// Get the URI for a cell by its index within the cell array.
-    #[expect(dead_code)]
     pub(crate) fn cell_uri_by_index(&self, index: CellId) -> Option<&lsp_types::Url> {
         self.cells.get(index).map(|cell| &cell.url)
     }
@@ -212,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

@@ -1,23 +1,51 @@
 use lsp_server::ErrorCode;
 use lsp_types::notification::PublishDiagnostics;
 use lsp_types::{
-    Diagnostic, DiagnosticRelatedInformation, DiagnosticSeverity, DiagnosticTag, NumberOrString,
-    PublishDiagnosticsParams, Range, Url,
+    CodeDescription, Diagnostic, DiagnosticRelatedInformation, DiagnosticSeverity, DiagnosticTag,
+    NumberOrString, PublishDiagnosticsParams, Range, Url,
 };
+use rustc_hash::FxHashMap;
 
 use ruff_db::diagnostic::{Annotation, Severity, SubDiagnostic};
 use ruff_db::files::FileRange;
 use ruff_db::source::{line_index, source_text};
 use ty_project::{Db, ProjectDatabase};
 
-use crate::DocumentSnapshot;
-use crate::PositionEncoding;
 use crate::document::{FileRangeExt, ToRangeExt};
 use crate::server::Result;
 use crate::server::client::Notifier;
+use crate::system::url_to_any_system_path;
+use crate::{DocumentSnapshot, PositionEncoding, Session};
 
 use super::LSPResult;
 
+/// 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`.
+///
+/// This is done by notifying the client with an empty list of diagnostics for the document.
 pub(super) fn clear_diagnostics(uri: &Url, notifier: &Notifier) -> Result<()> {
     notifier
         .notify::<PublishDiagnostics>(PublishDiagnosticsParams {
@@ -29,25 +57,106 @@ pub(super) fn clear_diagnostics(uri: &Url, notifier: &Notifier) -> Result<()> {
     Ok(())
 }
 
+/// Publishes the diagnostics for the given document snapshot using the [publish diagnostics
+/// notification].
+///
+/// This function is a no-op if the client supports pull diagnostics.
+///
+/// [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, notifier: &Notifier) -> Result<()> {
+    if session.client_capabilities().pull_diagnostics {
+        return Ok(());
+    }
+
+    let Ok(path) = url_to_any_system_path(&url) else {
+        return Ok(());
+    };
+
+    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 Some(diagnostics) = compute_diagnostics(db, &snapshot) else {
+        return Ok(());
+    };
+
+    // Sends a notification to the client with the diagnostics for the document.
+    let publish_diagnostics_notification = |uri: Url, diagnostics: Vec<Diagnostic>| {
+        notifier
+            .notify::<PublishDiagnostics>(PublishDiagnosticsParams {
+                uri,
+                diagnostics,
+                version: Some(snapshot.query().version()),
+            })
+            .with_failure_code(lsp_server::ErrorCode::InternalError)
+    };
+
+    match diagnostics {
+        Diagnostics::TextDocument(diagnostics) => {
+            publish_diagnostics_notification(url, diagnostics)?;
+        }
+        Diagnostics::NotebookDocument(cell_diagnostics) => {
+            for (cell_url, diagnostics) in cell_diagnostics {
+                publish_diagnostics_notification(cell_url, diagnostics)?;
+            }
+        }
+    }
+
+    Ok(())
+}
+
 pub(super) fn compute_diagnostics(
     db: &ProjectDatabase,
     snapshot: &DocumentSnapshot,
-) -> Vec<Diagnostic> {
+) -> Option<Diagnostics> {
     let Some(file) = snapshot.file(db) else {
         tracing::info!(
             "No file found for snapshot for `{}`",
             snapshot.query().file_url()
         );
-        return vec![];
+        return None;
     };
 
     let diagnostics = db.check_file(file);
 
-    diagnostics
-        .as_slice()
-        .iter()
-        .map(|message| to_lsp_diagnostic(db, message, snapshot.encoding()))
-        .collect()
+    if let Some(notebook) = snapshot.query().as_notebook() {
+        let mut cell_diagnostics: FxHashMap<Url, Vec<Diagnostic>> = FxHashMap::default();
+
+        // Populates all relevant URLs with an empty diagnostic list. This ensures that documents
+        // without diagnostics still get updated.
+        for cell_url in notebook.cell_urls() {
+            cell_diagnostics.entry(cell_url.clone()).or_default();
+        }
+
+        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;
+            };
+            cell_diagnostics
+                .entry(cell_uri.clone())
+                .or_default()
+                .push(diagnostic);
+        }
+
+        Some(Diagnostics::NotebookDocument(cell_diagnostics))
+    } else {
+        Some(Diagnostics::TextDocument(
+            diagnostics
+                .iter()
+                .map(|diagnostic| to_lsp_diagnostic(db, diagnostic, snapshot.encoding()))
+                .collect(),
+        ))
+    }
 }
 
 /// Converts the tool specific [`Diagnostic`][ruff_db::diagnostic::Diagnostic] to an LSP
@@ -91,9 +200,8 @@ fn to_lsp_diagnostic(
         .id()
         .is_lint()
         .then(|| {
-            Some(lsp_types::CodeDescription {
-                href: lsp_types::Url::parse(&format!("https://ty.dev/rules#{}", diagnostic.id()))
-                    .ok()?,
+            Some(CodeDescription {
+                href: Url::parse(&format!("https://ty.dev/rules#{}", diagnostic.id())).ok()?,
             })
         })
         .flatten();

crates/ty_server/src/server/api/notifications/did_change.rs

@@ -1,11 +1,12 @@
 use lsp_server::ErrorCode;
-use lsp_types::DidChangeTextDocumentParams;
 use lsp_types::notification::DidChangeTextDocument;
+use lsp_types::{DidChangeTextDocumentParams, VersionedTextDocumentIdentifier};
 
 use ty_project::watch::ChangeEvent;
 
 use crate::server::Result;
 use crate::server::api::LSPResult;
+use crate::server::api::diagnostics::publish_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::server::client::{Notifier, Requester};
 use crate::session::Session;
@@ -20,18 +21,23 @@ impl NotificationHandler for DidChangeTextDocumentHandler {
 impl SyncNotificationHandler for DidChangeTextDocumentHandler {
     fn run(
         session: &mut Session,
-        _notifier: Notifier,
+        notifier: Notifier,
         _requester: &mut Requester,
         params: DidChangeTextDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = url_to_any_system_path(&params.text_document.uri) else {
+        let DidChangeTextDocumentParams {
+            text_document: VersionedTextDocumentIdentifier { uri, version },
+            content_changes,
+        } = params;
+
+        let Ok(path) = url_to_any_system_path(&uri) else {
             return Ok(());
         };
 
-        let key = session.key_from_url(params.text_document.uri);
+        let key = session.key_from_url(uri.clone());
 
         session
-            .update_text_document(&key, params.content_changes, params.text_document.version)
+            .update_text_document(&key, content_changes, version)
             .with_failure_code(ErrorCode::InternalError)?;
 
         match path {
@@ -48,8 +54,6 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             }
         }
 
-        // TODO(dhruvmanila): Publish diagnostics if the client doesn't support pull diagnostics
-
-        Ok(())
+        publish_diagnostics(session, uri, &notifier)
     }
 }