Support pulling diagnostics for notebooks too

client-node-tests/src/integration.test.ts

@@ -8,6 +8,7 @@ import * as assert from 'assert';
 import * as path from 'path';
 import * as vscode from 'vscode';
 import * as lsclient from 'vscode-languageclient/node';
+import * as proto from 'vscode-languageserver-protocol';
 import { MemoryFileSystemProvider } from './memoryFileSystemProvider';
 import { vsdiag, DiagnosticProviderMiddleware } from 'vscode-languageclient/lib/common/diagnostic';
 
@@ -16,6 +17,11 @@ namespace GotNotifiedRequest {
 	export const type = new lsclient.RequestType<string, boolean, void>(method);
 }
 
+namespace SetDiagnosticsNotification {
+	export const method: 'testing/setDiagnostics' = 'testing/setDiagnostics';
+	export const type = new lsclient.NotificationType<proto.DocumentDiagnosticReport>(method);
+}
+
 async function revertAllDirty(): Promise<void> {
 	return vscode.commands.executeCommand('_workbench.revertAllDirty');
 }
@@ -1688,6 +1694,47 @@ suite('Full notebook tests', () => {
 		assert.strictEqual(notified, true);
 		await revertAllDirty();
 	});
+
+	test('Notebook document: pull diagnostics', async(): Promise<void> => {
+		const notebook = await vscode.workspace.openNotebookDocument('jupyter-notebook', createNotebookData());
+
+		// Send the diagnostics for the first cell.
+		const report: proto.DocumentDiagnosticReport = {
+			kind: 'full',
+			items: [
+				{ message: 'notebook-error', range: { start: { line: 0, character: 0 }, end: { line: 0, character: 0 }} }
+			]
+		};
+		await client.sendNotification(SetDiagnosticsNotification.method, { uri: notebook.cellAt(0).document.uri.toString(), report });
+
+		// notebook has to be visible for diagnostics to be published.
+		await vscode.window.showNotebookDocument(notebook);
+
+		const promise = new Promise<void>((resolve) => {
+			client.middleware.provideDiagnostics = async (doc, p, token, next) => {
+				const result = await next(doc, p, token);
+				if (result?.kind === 'full' && result?.items.length > 0) {
+					// Need to be async so that the diagnostics are published.
+					setTimeout(() => resolve(), 10);
+				}
+				return result;
+			};
+		});
+
+		// Change the notebook cell, this should cause the diagnostics to be published.
+		const edit = new vscode.WorkspaceEdit();
+		edit.insert(notebook.cellAt(0).document.uri, new vscode.Position(0, 0), '# a comment\n');
+		await vscode.workspace.applyEdit(edit);
+
+		// Wait for the diagnostics to be published.
+		await promise;
+		const diagnostics = vscode.languages.getDiagnostics(notebook.cellAt(0).document.uri);
+		assert.strictEqual(diagnostics.length, 1);
+		const diagnostic = diagnostics[0];
+		assert.strictEqual(diagnostic.message, 'notebook-error');
+
+		await revertAllDirty();
+	});
 });
 
 suite('Simple notebook tests', () => {

client-node-tests/src/servers/fullNotebookServer.ts

@@ -6,7 +6,10 @@
 import {
 	createConnection, InitializeParams, ServerCapabilities, TextDocumentSyncKind, RequestType,
 	DidOpenNotebookDocumentNotification, DidChangeNotebookDocumentNotification, DidSaveNotebookDocumentNotification,
-	DidCloseNotebookDocumentNotification
+	DidCloseNotebookDocumentNotification,
+	PublishDiagnosticsNotification,
+	DocumentDiagnosticReport,
+	NotificationType
 } from 'vscode-languageserver/node';
 
 const connection = createConnection();
@@ -20,6 +23,13 @@ namespace GotNotifiedRequest {
 	export const type = new RequestType<string, boolean, void>(method);
 }
 
+namespace SetDiagnosticsNotification {
+	export const method: 'testing/setDiagnostics' = 'testing/setDiagnostics';
+	export const type = new NotificationType<DocumentDiagnosticReport>(method);
+}
+
+const diagnostics = new Map<string, DocumentDiagnosticReport>();
+
 connection.onInitialize((_params: InitializeParams): any => {
 	const capabilities: ServerCapabilities = {
 		textDocumentSync: TextDocumentSyncKind.Full,
@@ -28,6 +38,12 @@ connection.onInitialize((_params: InitializeParams): any => {
 				notebook: { notebookType: 'jupyter-notebook' },
 				cells: [{ language: 'python' }]
 			}]
+		},
+		diagnosticProvider: {
+			identifier: 'diagnostic-provider',
+			documentSelector: null,
+			interFileDependencies: false,
+			workspaceDiagnostics: false
 		}
 	};
 	return { capabilities };
@@ -45,6 +61,11 @@ connection.notebooks.synchronization.onDidSaveNotebookDocument(() => {
 connection.notebooks.synchronization.onDidCloseNotebookDocument(() => {
 	receivedNotifications.add(DidCloseNotebookDocumentNotification.method);
 });
+connection.languages.diagnostics.on((params) => {
+	receivedNotifications.add(PublishDiagnosticsNotification.method);
+	const result = diagnostics.get(params.textDocument.uri) || { kind: 'unchanged', resultId: params.previousResultId || '' };
+	return result;
+});
 
 connection.onRequest(GotNotifiedRequest.type, (method: string) => {
 	const result = receivedNotifications.has(method);
@@ -54,5 +75,9 @@ connection.onRequest(GotNotifiedRequest.type, (method: string) => {
 	return result;
 });
 
+connection.onNotification(SetDiagnosticsNotification.method, (params) => {
+	diagnostics.set(params.uri, params.report);
+});
+
 // Listen on the connection
 connection.listen();

client/src/common/diagnostic.ts

@@ -8,20 +8,22 @@ import * as minimatch from 'minimatch';
 import {
 	Disposable, languages as Languages, window as Window, workspace as Workspace, CancellationToken, ProviderResult, Diagnostic as VDiagnostic,
 	CancellationTokenSource, TextDocument, CancellationError, Event as VEvent, EventEmitter, DiagnosticCollection, Uri, TabInputText, TabInputTextDiff,
-	TabChangeEvent, Event, TabInputCustom, workspace
+	TabChangeEvent, Event, TabInputCustom, workspace, TabInputNotebook,	NotebookCell
 } from 'vscode';
 
 import {
 	ClientCapabilities, ServerCapabilities, DocumentSelector, DidOpenTextDocumentNotification, DidChangeTextDocumentNotification,
 	DidSaveTextDocumentNotification, DidCloseTextDocumentNotification, LinkedMap, Touch, RAL, TextDocumentFilter, PreviousResultId,
 	DiagnosticRegistrationOptions, DiagnosticServerCancellationData, DocumentDiagnosticParams, DocumentDiagnosticRequest, DocumentDiagnosticReportKind,
-	WorkspaceDocumentDiagnosticReport, WorkspaceDiagnosticRequest, WorkspaceDiagnosticParams, DiagnosticOptions, DiagnosticRefreshRequest, DiagnosticTag
+	WorkspaceDocumentDiagnosticReport, WorkspaceDiagnosticRequest, WorkspaceDiagnosticParams, DiagnosticOptions, DiagnosticRefreshRequest, DiagnosticTag,
+	NotebookDocumentSyncRegistrationType
 } from 'vscode-languageserver-protocol';
 
 import { generateUuid } from './utils/uuid';
 import {
 	TextDocumentLanguageFeature, FeatureClient, LSPCancellationError
 } from './features';
+import { NotebookDocumentSyncFeature } from './notebook';
 
 function ensure<T, K extends keyof T>(target: T, key: K): T[K] {
 	if (target[key] === void 0) {
@@ -268,6 +270,16 @@ class Tabs {
 
 	public isVisible(document: TextDocument | Uri): boolean {
 		const uri = document instanceof Uri ? document : document.uri;
+		if (uri.scheme === NotebookDocumentSyncFeature.CellScheme) {
+			// Notebook cells aren't in the list of tabs, but the notebook should be.
+			return Workspace.notebookDocuments.some(notebook => {
+				if (this.open.has(notebook.uri.toString())) {
+					const cell = notebook.getCells().find(cell => cell.document.uri.toString() === uri.toString());
+					return cell !== undefined;
+				}
+				return false;
+			});
+		}
 		return this.open.has(uri.toString());
 	}
 
@@ -289,6 +301,8 @@ class Tabs {
 					uri = input.modified;
 				} else if (input instanceof TabInputCustom) {
 					uri = input.uri;
+				} else if (input instanceof TabInputNotebook) {
+					uri = input.uri;
 				}
 				if (uri !== undefined && !seen.has(uri.toString())) {
 					seen.add(uri.toString());
@@ -916,6 +930,11 @@ class DiagnosticFeatureProviderImpl implements DiagnosticProviderShape {
 				: Languages.match(documentSelector, document) > 0 && tabs.isVisible(document);
 		};
 
+		const matchesCell = (cell: NotebookCell): boolean => {
+			// Cells match if the language is allowed and if the containing notebook is visible.
+			return Languages.match(documentSelector, cell.document) > 0 && tabs.isVisible(cell.notebook.uri);
+		};
+
 		const isActiveDocument = (document: TextDocument | Uri): boolean => {
 			return document instanceof Uri
 				? this.activeTextDocument?.uri.toString() === document.toString()
@@ -970,6 +989,15 @@ class DiagnosticFeatureProviderImpl implements DiagnosticProviderShape {
 				this.diagnosticRequestor.pull(textDocument, () => { addToBackgroundIfNeeded(textDocument); });
 			}
 		}));
+		const notebookFeature = client.getFeature(NotebookDocumentSyncRegistrationType.method);
+		disposables.push(notebookFeature.onOpenNotificationSent((event) => {
+			// Send a pull for all opened cells in the notebook.
+			for (const cell of event.getCells()) {
+				if (matchesCell(cell)) {
+					this.diagnosticRequestor.pull(cell.document, () => { addToBackgroundIfNeeded(cell.document); });
+				}
+			}
+		}));
 
 		disposables.push(tabs.onOpen((opened) => {
 			for (const resource of opened) {
@@ -1007,6 +1035,15 @@ class DiagnosticFeatureProviderImpl implements DiagnosticProviderShape {
 				pulledTextDocuments.add(textDocument.uri.toString());
 			}
 		}
+		// Do the same for already open notebook cells.
+		for (const notebookDocument of Workspace.notebookDocuments) {
+			for (const cell of notebookDocument.getCells()) {
+				if (matchesCell(cell)) {
+					this.diagnosticRequestor.pull(cell.document, () => { addToBackgroundIfNeeded(cell.document); });
+					pulledTextDocuments.add(cell.document.uri.toString());
+				}
+			}
+		}
 
 		// Pull all tabs if not already pulled as text document
 		if (diagnosticPullOptions.onTabs === true) {
@@ -1029,6 +1066,31 @@ class DiagnosticFeatureProviderImpl implements DiagnosticProviderShape {
 					this.diagnosticRequestor.pull(textDocument, () => { this.backgroundScheduler.trigger(); });
 				}
 			}));
+			disposables.push(notebookFeature.onChangeNotificationSent(async (event) => {
+				// Send a pull for all changed cells in the notebook.
+				const textEvents = event.cells?.textContent || [];
+				const changedCells = textEvents.map(
+					(c) => event.notebook.getCells().find(cell => cell.document.uri.toString() === c.document.uri.toString()));
+				for (const cell of changedCells) {
+					if (cell && matchesCell(cell)) {
+						this.diagnosticRequestor.pull(cell.document, () => { this.backgroundScheduler.trigger(); });
+					}
+				}
+
+				// Clear out any closed cells.
+				const closedCells = event.cells?.structure?.didClose || [];
+				for (const cell of closedCells) {
+					this.diagnosticRequestor.forgetDocument(cell.document);
+				}
+
+				// Send a pull for any new opened cells.
+				const openedCells = event.cells?.structure?.didOpen || [];
+				for (const cell of openedCells) {
+					if (matchesCell(cell)) {
+						this.diagnosticRequestor.pull(cell.document, () => { this.backgroundScheduler.trigger(); });
+					}
+				}
+			}));
 		}
 
 		if (diagnosticPullOptions.onSave === true) {
@@ -1039,13 +1101,25 @@ class DiagnosticFeatureProviderImpl implements DiagnosticProviderShape {
 					this.diagnosticRequestor.pull(event.textDocument);
 				}
 			}));
+			disposables.push(notebookFeature.onSaveNotificationSent((event) => {
+				for (const cell of event.getCells()) {
+					if (matchesCell(cell)) {
+						this.diagnosticRequestor.pull(cell.document);
+					}
+				}
+			}));
 		}
 
 		// When the document closes clear things up
 		const closeFeature = client.getFeature(DidCloseTextDocumentNotification.method);
 		disposables.push(closeFeature.onNotificationSent((event) => {
 			this.cleanUpDocument(event.textDocument);
 		}));
+		disposables.push(notebookFeature.onCloseNotificationSent((event) => {
+			for (const cell of event.getCells()) {
+				this.cleanUpDocument(cell.document);
+			}
+		}));
 
 		// Same when a tabs closes.
 		tabs.onClose((closed) => {