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/package.json

@@ -9,7 +9,8 @@
 		"vscode": "^1.57.0"
 	},
 	"enabledApiProposals": [
-    "notebookWorkspaceEdit"
+    "notebookWorkspaceEdit",
+    "documentPaste"
 	],
 	"activationEvents": [
 		"*"
@@ -27,6 +28,21 @@
 		}
 	},
 	"contributes": {
+    "configuration":[
+      {
+        "properties": {
+          "ipynb.experimental.pasteImages.enabled":{
+            "type": "boolean",
+            "scope": "resource",
+            "markdownDescription": "%ipynb.experimental.pasteImages.enabled%",
+            "default": false,
+            "tags": [
+              "experimental"
+            ]
+          }
+        }
+      }
+    ],
 		"commands": [
 			{
 				"command": "ipynb.newUntitledIpynb",

extensions/ipynb/package.nls.json

@@ -1,4 +1,5 @@
 {
 	"displayName": ".ipynb support",
-	"description": "Provides basic support for opening and reading Jupyter's .ipynb notebook files"
+	"description": "Provides basic support for opening and reading Jupyter's .ipynb notebook files",
+	"ipynb.experimental.pasteImages.enabled":"Enable/Disable pasting images into markdown cells within ipynb files. Requires enabling `#ipynb.experimental.pasteImages.enabled#`."
 }

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;
+}

extensions/ipynb/tsconfig.json

@@ -9,6 +9,7 @@
 	"include": [
 		"src/**/*",
 		"../../src/vscode-dts/vscode.d.ts",
-		"../../src/vscode-dts/vscode.proposed.notebookWorkspaceEdit.d.ts"
+		"../../src/vscode-dts/vscode.proposed.notebookWorkspaceEdit.d.ts",
+		"../../src/vscode-dts/vscode.proposed.documentPaste.d.ts"
 	]
 }

src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts

@@ -13,6 +13,7 @@ import { getExtensionForMimeType } from 'vs/base/common/mime';
 import { FileAccess, Schemas } from 'vs/base/common/network';
 import { isMacintosh, isWeb } from 'vs/base/common/platform';
 import { dirname, joinPath } from 'vs/base/common/resources';
+import { equals } from 'vs/base/common/objects';
 import { URI } from 'vs/base/common/uri';
 import * as UUID from 'vs/base/common/uuid';
 import { TokenizationRegistry } from 'vs/editor/common/languages';
@@ -1042,7 +1043,7 @@ var requirejs = (function() {
 		}
 
 		const sameContent = newContent.content === entry.content;
-		const sameMetadata = newContent.metadata === entry.metadata;
+		const sameMetadata = (equals(newContent.metadata, entry.metadata));
 		if (!sameContent || !sameMetadata || !entry.visible) {
 			this._sendMessageToWebview({
 				type: 'showMarkupCell',

src/vs/workbench/contrib/notebook/browser/view/renderers/webviewPreloads.ts

@@ -1671,6 +1671,7 @@ async function webviewPreloads(ctx: PreloadContext) {
 		private _content: { readonly value: string; readonly version: number; readonly metadata: NotebookCellMetadata };
 
 		constructor(id: string, mime: string, content: string, top: number, metadata: NotebookCellMetadata) {
+			const self = this;
 			this.id = id;
 			this._content = { value: content, version: 0, metadata: metadata };
 
@@ -1682,8 +1683,8 @@ async function webviewPreloads(ctx: PreloadContext) {
 				id,
 				mime,
 
-				metadata: (): NotebookCellMetadata => {
-					return this._content.metadata;
+				get metadata(): NotebookCellMetadata {
+					return self._content.metadata;
 				},
 
 				text: (): string => {