Add support for attachment cleaning on notebook save

extensions/ipynb/src/helper.ts

@@ -77,61 +77,66 @@ export function objectEquals(one: any, other: any) {
 	return true;
 }
 
-interface Options<T> {
-	callback: (value: T) => void;
-
-	merge?: (input: T[]) => T;
-	delay?: number;
-}
-
-
-export class DebounceTrigger<T> {
-
-	private _isPaused = 0;
-	protected _queue: T[] = [];
-	private _callbackFn: (value: T) => void;
-	private _mergeFn?: (input: T[]) => T;
-	private readonly _delay: number;
-	private _handle: any | undefined;
-
-	constructor(options: Options<T>) {
-		this._callbackFn = options.callback;
-		this._mergeFn = options.merge;
-		this._delay = options.delay ?? 100;
+/**
+ * A helper to delay/debounce execution of a task, includes cancellation/disposal support.
+ * Pulled from https://github.com/microsoft/vscode/blob/3059063b805ed0ac10a6d9539e213386bfcfb852/extensions/markdown-language-features/src/util/async.ts
+ */
+export class Delayer<T> {
+
+	public defaultDelay: number;
+	private _timeout: any; // Timer
+	private _cancelTimeout: Promise<T | null> | null;
+	private _onSuccess: ((value: T | PromiseLike<T> | undefined) => void) | null;
+	private _task: ITask<T> | null;
+
+	constructor(defaultDelay: number) {
+		this.defaultDelay = defaultDelay;
+		this._timeout = null;
+		this._cancelTimeout = null;
+		this._onSuccess = null;
+		this._task = null;
 	}
 
-	private pause(): void {
-		this._isPaused++;
+	dispose() {
+		this._doCancelTimeout();
 	}
 
-	private resume(): void {
-		if (this._isPaused !== 0 && --this._isPaused === 0) {
-			if (this._mergeFn) {
-				const items = Array.from(this._queue);
-				this._queue = [];
-				this._callbackFn(this._mergeFn(items));
-
-			} else {
-				while (!this._isPaused && this._queue.length !== 0) {
-					this._callbackFn(this._queue.shift()!);
-				}
-			}
+	public trigger(task: ITask<T>, delay: number = this.defaultDelay): Promise<T | null> {
+		this._task = task;
+		if (delay >= 0) {
+			this._doCancelTimeout();
 		}
-	}
 
-	trigger(item: T): void {
-		if (!this._handle) {
-			this.pause();
-			this._handle = setTimeout(() => {
-				this._handle = undefined;
-				this.resume();
-			}, this._delay);
+		if (!this._cancelTimeout) {
+			this._cancelTimeout = new Promise<T | undefined>((resolve) => {
+				this._onSuccess = resolve;
+			}).then(() => {
+				this._cancelTimeout = null;
+				this._onSuccess = null;
+				const result = this._task && this._task?.();
+				this._task = null;
+				return result;
+			});
 		}
 
-		if (this._isPaused !== 0) {
-			this._queue.push(item);
-		} else {
-			this._callbackFn(item);
+		if (delay >= 0 || this._timeout === null) {
+			this._timeout = setTimeout(() => {
+				this._timeout = null;
+				this._onSuccess?.(undefined);
+			}, delay >= 0 ? delay : this.defaultDelay);
 		}
+
+		return this._cancelTimeout;
 	}
+
+	private _doCancelTimeout(): void {
+		if (this._timeout !== null) {
+			clearTimeout(this._timeout);
+			this._timeout = null;
+		}
+	}
+}
+
+export interface ITask<T> {
+	(): T;
 }

extensions/ipynb/src/notebookAttachmentCleaner.ts

@@ -5,7 +5,7 @@
 
 import * as vscode from 'vscode';
 import { ATTACHMENT_CLEANUP_COMMANDID, JUPYTER_NOTEBOOK_MARKDOWN_SELECTOR } from './constants';
-import { DebounceTrigger, deepClone, objectEquals } from './helper';
+import { deepClone, objectEquals, Delayer } from './helper';
 
 interface AttachmentCleanRequest {
 	notebook: vscode.NotebookDocument;
@@ -32,15 +32,10 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
 
 	private _disposables: vscode.Disposable[];
 	private _imageDiagnosticCollection: vscode.DiagnosticCollection;
+	private readonly _delayer = new Delayer(750);
+
 	constructor() {
 		this._disposables = [];
-		const debounceTrigger = new DebounceTrigger<AttachmentCleanRequest>({
-			callback: (change: AttachmentCleanRequest) => {
-				this.cleanNotebookAttachments(change);
-			},
-			delay: 500
-		});
-
 		this._imageDiagnosticCollection = vscode.languages.createDiagnosticCollection('Notebook Image Attachment');
 		this._disposables.push(this._imageDiagnosticCollection);
 
@@ -57,23 +52,66 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
 		}));
 
 		this._disposables.push(vscode.workspace.onDidChangeNotebookDocument(e => {
-			e.cellChanges.forEach(change => {
-				if (!change.document) {
-					return;
-				}
+			this._delayer.trigger(() => {
 
-				if (change.cell.kind !== vscode.NotebookCellKind.Markup) {
-					return;
-				}
+				e.cellChanges.forEach(change => {
+					if (!change.document) {
+						return;
+					}
 
-				debounceTrigger.trigger({
-					notebook: e.notebook,
-					cell: change.cell,
-					document: change.document
+					if (change.cell.kind !== vscode.NotebookCellKind.Markup) {
+						return;
+					}
+
+					const metadataEdit = this.cleanNotebookAttachments({
+						notebook: e.notebook,
+						cell: change.cell,
+						document: change.document
+					});
+					if (metadataEdit) {
+						const workspaceEdit = new vscode.WorkspaceEdit();
+						workspaceEdit.set(e.notebook.uri, [metadataEdit]);
+						vscode.workspace.applyEdit(workspaceEdit);
+					}
 				});
 			});
 		}));
 
+
+		this._disposables.push(vscode.workspace.onWillSaveNotebookDocument(e => {
+			if (e.reason === vscode.TextDocumentSaveReason.Manual) {
+				this._delayer.dispose();
+
+				e.waitUntil(new Promise((resolve) => {
+					if (e.notebook.getCells().length === 0) {
+						return;
+					}
+
+					const notebookEdits: vscode.NotebookEdit[] = [];
+					for (const cell of e.notebook.getCells()) {
+						if (cell.kind !== vscode.NotebookCellKind.Markup) {
+							continue;
+						}
+
+						const metadataEdit = this.cleanNotebookAttachments({
+							notebook: e.notebook,
+							cell: cell,
+							document: cell.document
+						});
+
+						if (metadataEdit) {
+							notebookEdits.push(metadataEdit);
+						}
+					}
+
+					const workspaceEdit = new vscode.WorkspaceEdit();
+					workspaceEdit.set(e.notebook.uri, notebookEdits);
+
+					resolve(workspaceEdit);
+				}));
+			}
+		}));
+
 		this._disposables.push(vscode.workspace.onDidCloseNotebookDocument(e => {
 			this._attachmentCache.delete(e.uri.toString());
 		}));
@@ -134,8 +172,10 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
 	/**
 	 * take in a NotebookDocumentChangeEvent, and clean the attachment data for the cell(s) that have had their markdown source code changed
 	 * @param e NotebookDocumentChangeEvent from the onDidChangeNotebookDocument listener
+	 * @returns vscode.NotebookEdit, the metadata alteration performed on the json behind the ipynb
 	 */
-	private cleanNotebookAttachments(e: AttachmentCleanRequest) {
+	private cleanNotebookAttachments(e: AttachmentCleanRequest): vscode.NotebookEdit | undefined {
+
 		if (e.notebook.isClosed) {
 			return;
 		}
@@ -187,16 +227,16 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
 			}
 		}
 
+		this.updateDiagnostics(cell.document.uri, diagnostics);
+
 		if (cell.index > -1 && !objectEquals(markdownAttachmentsInUse, cell.metadata.attachments)) {
 			const updateMetadata: { [key: string]: any } = deepClone(cell.metadata);
 			updateMetadata.attachments = markdownAttachmentsInUse;
 			const metadataEdit = vscode.NotebookEdit.updateCellMetadata(cell.index, updateMetadata);
-			const workspaceEdit = new vscode.WorkspaceEdit();
-			workspaceEdit.set(e.notebook.uri, [metadataEdit]);
-			vscode.workspace.applyEdit(workspaceEdit);
-		}
 
-		this.updateDiagnostics(cell.document.uri, diagnostics);
+			return metadataEdit;
+		}
+		return;
 	}
 
 	private analyzeMissingAttachments(document: vscode.TextDocument): void {
@@ -345,6 +385,7 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
 
 	dispose() {
 		this._disposables.forEach(d => d.dispose());
+		this._delayer.dispose();
 	}
 }
 

extensions/ipynb/tsconfig.json

@@ -2,9 +2,7 @@
 	"extends": "../tsconfig.base.json",
 	"compilerOptions": {
 		"outDir": "./out",
-		"lib": [
-			"dom"
-		]
+		"lib": ["dom"]
 	},
 	"include": [
 		"src/**/*",