Support for pasting images into markdown notebook cells (#156847)

* dataflow support for updated metadata

* update cellAttachmentRenderer.ts to reflect metadata being a getter() inside MarkupCell

* document paste additions

* update condition to re-render cells, now includes metadata changes

* paste API working, debugging command added

* paste working with metadata. needs numbering, and cleaning upon delete

* paste screenshot works fully

* remove debugging command. Cleaning.

* notebook cells now re-render upon metadata changes

* changed name validity checking, remove unneeded function

* use _document for cell data, use snippet choice, dto fix

* return subscription, for loop, uri fix, alter metadata in-place, better snippet

* metadata fix, object.equals, fix cellAttRenderer metadata call

* added comment with source of encodeBase64

* gate mkdn image paste behind experimental setting

extensions/ipynb/src/cellAttachmentRenderer.ts

@@ -22,7 +22,7 @@ export async function activate(ctx: RendererContext<void>) {
 		md.renderer.rules.image = (tokens: MarkdownItToken[], idx: number, options, env, self) => {
 			const token = tokens[idx];
 			const src = token.attrGet('src');
-			const attachments: Record<string, Record<string, string>> = env.outputItem.metadata().custom?.attachments;
+			const attachments: Record<string, Record<string, string>> = env.outputItem.metadata.custom?.attachments;
 			if (attachments && src) {
 				const imageAttachment = attachments[src.replace('attachment:', '')];
 				if (imageAttachment) {

extensions/ipynb/src/ipynbMain.ts

@@ -6,6 +6,7 @@
 import * as vscode from 'vscode';
 import { ensureAllNewCellsHaveCellIds } from './cellIdService';
 import { NotebookSerializer } from './notebookSerializer';
+import * as NotebookImagePaste from './notebookImagePaste';
 
 // From {nbformat.INotebookMetadata} in @jupyterlab/coreutils
 type NotebookMetadata = {
@@ -77,12 +78,15 @@ export function activate(context: vscode.ExtensionContext) {
 		await vscode.window.showNotebookDocument(document);
 	}));
 
+	context.subscriptions.push(NotebookImagePaste.imagePasteSetup());
+
 	// Update new file contribution
 	vscode.extensions.onDidChange(() => {
 		vscode.commands.executeCommand('setContext', 'jupyterEnabled', vscode.extensions.getExtension('ms-toolsai.jupyter'));
 	});
 	vscode.commands.executeCommand('setContext', 'jupyterEnabled', vscode.extensions.getExtension('ms-toolsai.jupyter'));
 
+
 	return {
 		exportNotebook: (notebook: vscode.NotebookData): string => {
 			return exportNotebook(notebook, serializer);

extensions/ipynb/src/notebookImagePaste.ts

@@ -0,0 +1,143 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+
+class CopyPasteEditProvider implements vscode.DocumentPasteEditProvider {
+
+	async provideDocumentPasteEdits(
+		_document: vscode.TextDocument,
+		_ranges: readonly vscode.Range[],
+		dataTransfer: vscode.DataTransfer,
+		_token: vscode.CancellationToken
+	): Promise<vscode.DocumentPasteEdit | undefined> {
+
+		const enabled = vscode.workspace.getConfiguration('ipynb', _document).get('experimental.pasteImages.enabled', false);
+		if (!enabled) {
+			return undefined;
+		}
+
+		// get b64 data from paste
+		// TODO: dataTransfer.get() limits to one image pasted
+		const dataItem = dataTransfer.get('image/png');
+		if (!dataItem) {
+			return undefined;
+		}
+		const fileDataAsUint8 = await dataItem.asFile()?.data();
+		if (!fileDataAsUint8) {
+			return undefined;
+		}
+
+		// get filename data from paste
+		let pasteFilename = dataItem.asFile()?.name;
+		if (!pasteFilename) {
+			return undefined;
+		}
+		const separatorIndex = pasteFilename?.lastIndexOf('.');
+		const filename = pasteFilename?.slice(0, separatorIndex);
+		const filetype = pasteFilename?.slice(separatorIndex);
+		if (!filename || !filetype) {
+			return undefined;
+		}
+
+		// get notebook cell data
+		let notebookUri;
+		let currentCell;
+		for (const notebook of vscode.workspace.notebookDocuments) {
+			if (notebook.uri.path === _document.uri.path) {
+				for (const cell of notebook.getCells()) {
+					if (cell.document === _document) {
+						currentCell = cell;
+						notebookUri = notebook.uri;
+						break;
+					}
+				}
+			}
+		}
+		if (!currentCell || !notebookUri) {
+			return undefined;
+		}
+
+		// create updated metadata for cell (prep for WorkspaceEdit)
+		const b64string = encodeBase64(fileDataAsUint8);
+		const startingAttachments = currentCell.metadata?.custom?.attachments;
+		if (!startingAttachments) {
+			currentCell.metadata.custom['attachments'] = { [pasteFilename]: { 'image/png': b64string } };
+		} else {
+			for (let appendValue = 2; pasteFilename in startingAttachments; appendValue++) {
+				const objEntries = Object.entries(startingAttachments[pasteFilename]);
+				if (objEntries.length) { // check that mime:b64 are present
+					const [, attachmentb64] = objEntries[0];
+					if (attachmentb64 !== b64string) {	// append a "-#" here. same name, diff data. this matches jupyter behavior
+						pasteFilename = filename.concat(`-${appendValue}`) + filetype;
+					}
+				}
+			}
+			currentCell.metadata.custom.attachments[pasteFilename] = { 'image/png': b64string };
+		}
+
+		const metadataNotebookEdit = vscode.NotebookEdit.updateCellMetadata(currentCell.index, currentCell.metadata);
+		const workspaceEdit = new vscode.WorkspaceEdit();
+		if (metadataNotebookEdit) {
+			workspaceEdit.set(notebookUri, [metadataNotebookEdit]);
+		}
+
+		// create a snippet for paste
+		const pasteSnippet = new vscode.SnippetString();
+		pasteSnippet.appendText('![');
+		pasteSnippet.appendPlaceholder(`${pasteFilename}`);
+		pasteSnippet.appendText(`](attachment:${pasteFilename})`);
+
+		return { insertText: pasteSnippet, additionalEdit: workspaceEdit };
+	}
+}
+
+export function imagePasteSetup() {
+	const selector: vscode.DocumentSelector = { notebookType: 'jupyter-notebook', language: 'markdown' }; // this is correct provider
+	return vscode.languages.registerDocumentPasteEditProvider(selector, new CopyPasteEditProvider(), {
+		pasteMimeTypes: ['image/png'],
+	});
+}
+
+/**
+ *  Taken from https://github.com/microsoft/vscode/blob/743b016722db90df977feecde0a4b3b4f58c2a4c/src/vs/base/common/buffer.ts#L350-L387
+ */
+function encodeBase64(buffer: Uint8Array, padded = true, urlSafe = false) {
+	const base64Alphabet = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/';
+	const base64UrlSafeAlphabet = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_';
+
+	const dictionary = urlSafe ? base64UrlSafeAlphabet : base64Alphabet;
+	let output = '';
+
+	const remainder = buffer.byteLength % 3;
+
+	let i = 0;
+	for (; i < buffer.byteLength - remainder; i += 3) {
+		const a = buffer[i + 0];
+		const b = buffer[i + 1];
+		const c = buffer[i + 2];
+
+		output += dictionary[a >>> 2];
+		output += dictionary[(a << 4 | b >>> 4) & 0b111111];
+		output += dictionary[(b << 2 | c >>> 6) & 0b111111];
+		output += dictionary[c & 0b111111];
+	}
+
+	if (remainder === 1) {
+		const a = buffer[i + 0];
+		output += dictionary[a >>> 2];
+		output += dictionary[(a << 4) & 0b111111];
+		if (padded) { output += '=='; }
+	} else if (remainder === 2) {
+		const a = buffer[i + 0];
+		const b = buffer[i + 1];
+		output += dictionary[a >>> 2];
+		output += dictionary[(a << 4 | b >>> 4) & 0b111111];
+		output += dictionary[(b << 2) & 0b111111];
+		if (padded) { output += '='; }
+	}
+
+	return output;
+}