[ty] Support pull diagnostics for notebook cells (#27778)

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

@@ -197,16 +197,12 @@ pub(super) enum LspDiagnostics {
 }
 
 impl LspDiagnostics {
-    /// 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> {
+    /// Returns the diagnostics for the text document or notebook cell at `uri`.
+    pub(super) fn into_document_diagnostics(self, uri: &Uri) -> Vec<Diagnostic> {
         match self {
             LspDiagnostics::TextDocument(diagnostics) => diagnostics,
-            LspDiagnostics::NotebookDocument(_) => {
-                panic!("Expected a text document diagnostics, but got notebook diagnostics")
+            LspDiagnostics::NotebookDocument(mut diagnostics) => {
+                diagnostics.remove(uri).unwrap_or_default()
             }
         }
     }

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

@@ -55,24 +55,22 @@ impl BackgroundDocumentRequestHandler for DocumentDiagnosticRequestHandler {
                 }
                 .into()
             }
-            new_id => {
-                RelatedFullDocumentDiagnosticReport {
-                    related_documents: None,
-                    full_document_diagnostic_report: FullDocumentDiagnosticReport {
-                        result_id: new_id,
-                        // SAFETY: Pull diagnostic requests are only called for text documents, not for
-                        // notebook documents.
-                        items: diagnostics
-                            .to_lsp_diagnostics(
-                                db,
-                                snapshot.resolved_client_capabilities(),
-                                snapshot.global_settings(),
-                            )
-                            .expect_text_document(),
-                    },
-                }
-                .into()
+            new_id => RelatedFullDocumentDiagnosticReport {
+                related_documents: None,
+                full_document_diagnostic_report: FullDocumentDiagnosticReport {
+                    result_id: new_id,
+                    // A notebook is checked as a whole, but a pull response only includes
+                    // diagnostics for the requested cell.
+                    items: diagnostics
+                        .to_lsp_diagnostics(
+                            db,
+                            snapshot.resolved_client_capabilities(),
+                            snapshot.global_settings(),
+                        )
+                        .into_document_diagnostics(snapshot.uri()),
+                },
             }
+            .into(),
         };
 
         Ok(report)

crates/ty_server/tests/e2e/notebook.rs

@@ -59,6 +59,52 @@ type Style = Literal["italic", "bold", "underline"]"#,
     Ok(())
 }
 
+#[test]
+fn pull_diagnostics_for_notebook_cells() -> anyhow::Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .enable_pull_diagnostics(true)
+        .build()
+        .wait_until_workspaces_are_initialized();
+
+    let mut builder = NotebookBuilder::virtual_file("test.ipynb");
+    let first_cell = builder.add_python_cell("value: str = 1\n");
+    let second_cell = builder.add_python_cell("def example():\n    unused = 1\n    return 0\n");
+    let third_cell = builder.add_python_cell("value.upper()\n");
+
+    builder.open(&mut server);
+    server.collect_publish_diagnostic_notifications(3);
+
+    let diagnostics = [first_cell, second_cell, third_cell]
+        .into_iter()
+        .map(|uri| {
+            let id = server.send_request::<lsp_types::DocumentDiagnosticRequest>(
+                lsp_types::DocumentDiagnosticParams {
+                    text_document: TextDocumentIdentifier { uri },
+                    identifier: Some("ty".to_string()),
+                    previous_result_id: None,
+                    work_done_progress_params: lsp_types::WorkDoneProgressParams::default(),
+                    partial_result_params: lsp_types::PartialResultParams::default(),
+                },
+            );
+
+            match server.await_response::<lsp_types::DocumentDiagnosticRequest>(&id) {
+                lsp_types::DocumentDiagnosticReport::RelatedFullDocumentDiagnosticReport(
+                    report,
+                ) => report.full_document_diagnostic_report.items,
+                lsp_types::DocumentDiagnosticReport::RelatedUnchangedDocumentDiagnosticReport(
+                    _,
+                ) => {
+                    panic!("Expected a full diagnostic report")
+                }
+            }
+        })
+        .collect::<Vec<_>>();
+
+    assert_json_snapshot!(diagnostics);
+
+    Ok(())
+}
+
 #[test]
 fn publish_unused_binding_diagnostics_open() -> anyhow::Result<()> {
     let mut server = TestServerBuilder::new()?

crates/ty_server/tests/e2e/snapshots/e2e__notebook__pull_diagnostics_for_notebook_cells.snap

@@ -0,0 +1,48 @@
+---
+source: crates/ty_server/tests/e2e/notebook.rs
+expression: diagnostics
+---
+[
+  [
+    {
+      "range": {
+        "start": {
+          "line": 0,
+          "character": 13
+        },
+        "end": {
+          "line": 0,
+          "character": 14
+        }
+      },
+      "severity": 1,
+      "code": "invalid-assignment",
+      "codeDescription": {
+        "href": "https://ty.dev/rules#invalid-assignment"
+      },
+      "source": "ty",
+      "message": "Object of type `Literal[1]` is not assignable to `str`"
+    }
+  ],
+  [
+    {
+      "range": {
+        "start": {
+          "line": 1,
+          "character": 4
+        },
+        "end": {
+          "line": 1,
+          "character": 10
+        }
+      },
+      "severity": 4,
+      "source": "ty",
+      "message": "`unused` is unused",
+      "tags": [
+        1
+      ]
+    }
+  ],
+  []
+]