Add cell_ids for ipynb with nbformat >= 4.5 (#134835)

extensions/ipynb/package.json

@@ -8,6 +8,7 @@
 	"engines": {
 		"vscode": "^1.57.0"
 	},
+  "enableProposedApi": true,
 	"activationEvents": [
 		"onNotebook:jupyter-notebook"
 	],
@@ -64,10 +65,12 @@
   },
 	"dependencies": {
 		"@enonic/fnv-plus": "^1.3.0",
-		"detect-indent": "^6.0.0"
+		"detect-indent": "^6.0.0",
+		"uuid": "^8.3.2"
 	},
 	"devDependencies": {
-		"@jupyterlab/coreutils": "^3.1.0"
+		"@jupyterlab/coreutils": "^3.1.0",
+		"@types/uuid": "^8.3.1"
 	},
   "repository": {
     "type": "git",

extensions/ipynb/src/cellIdService.ts

@@ -0,0 +1,77 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { ExtensionContext, NotebookCellsChangeEvent, NotebookDocument, notebooks, workspace, WorkspaceEdit } from 'vscode';
+import { v4 as uuid } from 'uuid';
+import { getCellMetadata } from './serializers';
+import { CellMetadata } from './common';
+import { getNotebookMetadata } from './notebookSerializer';
+import { nbformat } from '@jupyterlab/coreutils';
+
+/**
+ * Ensure all new cells in notebooks with nbformat >= 4.5 have an id.
+ * Details of the spec can be found here https://jupyter.org/enhancement-proposals/62-cell-id/cell-id.html#
+ */
+export function ensureAllNewCellsHaveCellIds(context: ExtensionContext) {
+	notebooks.onDidChangeNotebookCells(onDidChangeNotebookCells, undefined, context.subscriptions);
+}
+
+function onDidChangeNotebookCells(e: NotebookCellsChangeEvent) {
+	const nbMetadata = getNotebookMetadata(e.document);
+	if (!isCellIdRequired(nbMetadata)) {
+		return;
+	}
+	e.changes.forEach(change => {
+		change.items.forEach(cell => {
+			const cellMetadata = getCellMetadata(cell);
+			if (cellMetadata?.id) {
+				return;
+			}
+			const id = generateCellId(e.document);
+			const edit = new WorkspaceEdit();
+			// Don't edit the metadata directly, always get a clone (prevents accidental singletons and directly editing the objects).
+			const updatedMetadata: CellMetadata = { ...JSON.parse(JSON.stringify(cellMetadata || {})) };
+			updatedMetadata.id = id;
+			edit.replaceNotebookCellMetadata(cell.notebook.uri, cell.index, { ...(cell.metadata), custom: updatedMetadata });
+			workspace.applyEdit(edit);
+		});
+	});
+}
+
+/**
+ * Cell ids are required in notebooks only in notebooks with nbformat >= 4.5
+ */
+function isCellIdRequired(metadata: Pick<Partial<nbformat.INotebookContent>, 'nbformat' | 'nbformat_minor'>) {
+	if ((metadata.nbformat || 0) >= 5) {
+		return true;
+	}
+	if ((metadata.nbformat || 0) === 4 && (metadata.nbformat_minor || 0) >= 5) {
+		return true;
+	}
+	return false;
+}
+
+function generateCellId(notebook: NotebookDocument) {
+	while (true) {
+		// Details of the id can be found here https://jupyter.org/enhancement-proposals/62-cell-id/cell-id.html#adding-an-id-field,
+		// & here https://jupyter.org/enhancement-proposals/62-cell-id/cell-id.html#updating-older-formats
+		const id = uuid().replace(/-/g, '').substring(0, 8);
+		let duplicate = false;
+		for (let index = 0; index < notebook.cellCount; index++) {
+			const cell = notebook.cellAt(index);
+			const existingId = getCellMetadata(cell)?.id;
+			if (!existingId) {
+				continue;
+			}
+			if (existingId === id) {
+				duplicate = true;
+				break;
+			}
+		}
+		if (!duplicate) {
+			return id;
+		}
+	}
+}

extensions/ipynb/src/ipynbMain.ts

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as vscode from 'vscode';
+import { ensureAllNewCellsHaveCellIds } from './cellIdService';
 import { NotebookSerializer } from './notebookSerializer';
 
 // From {nbformat.INotebookMetadata} in @jupyterlab/coreutils
@@ -27,6 +28,7 @@ type NotebookMetadata = {
 
 export function activate(context: vscode.ExtensionContext) {
 	const serializer = new NotebookSerializer(context);
+	ensureAllNewCellsHaveCellIds(context);
 	context.subscriptions.push(vscode.workspace.registerNotebookSerializer('jupyter-notebook', serializer, {
 		transientOutputs: false,
 		transientCellMetadata: {

extensions/ipynb/src/notebookSerializer.ts

@@ -78,11 +78,7 @@ export class NotebookSerializer implements vscode.NotebookSerializer {
 	}
 
 	public serializeNotebookToString(data: vscode.NotebookData): string {
-		const notebookContent: Partial<nbformat.INotebookContent> = data.metadata?.custom || {};
-		notebookContent.cells = notebookContent.cells || [];
-		notebookContent.nbformat = notebookContent.nbformat || 4;
-		notebookContent.nbformat_minor = notebookContent.nbformat_minor ?? 2;
-		notebookContent.metadata = notebookContent.metadata || { orig_nbformat: 4 };
+		const notebookContent = getNotebookMetadata(data);
 
 		notebookContent.cells = data.cells
 			.map(cell => createJupyterCellFromNotebookCell(cell))
@@ -95,3 +91,12 @@ export class NotebookSerializer implements vscode.NotebookSerializer {
 		return JSON.stringify(sortObjectPropertiesRecursively(notebookContent), undefined, indentAmount) + '\n';
 	}
 }
+
+export function getNotebookMetadata(document: vscode.NotebookDocument | vscode.NotebookData) {
+	const notebookContent: Partial<nbformat.INotebookContent> = document.metadata?.custom || {};
+	notebookContent.cells = notebookContent.cells || [];
+	notebookContent.nbformat = notebookContent.nbformat || 4;
+	notebookContent.nbformat_minor = notebookContent.nbformat_minor ?? 2;
+	notebookContent.metadata = notebookContent.metadata || { orig_nbformat: 4 };
+	return notebookContent;
+}

extensions/ipynb/src/serializers.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { nbformat } from '@jupyterlab/coreutils';
-import { NotebookCellData, NotebookCellKind, NotebookCellOutput } from 'vscode';
+import { NotebookCell, NotebookCellData, NotebookCellKind, NotebookCellOutput } from 'vscode';
 import { CellMetadata, CellOutputMetadata } from './common';
 import { textMimeTypes } from './deserializers';
 
@@ -53,8 +53,11 @@ export function sortObjectPropertiesRecursively(obj: any): any {
 	return obj;
 }
 
+export function getCellMetadata(cell: NotebookCell | NotebookCellData) {
+	return cell.metadata?.custom as CellMetadata | undefined;
+}
 function createCodeCellFromNotebookCell(cell: NotebookCellData): nbformat.ICodeCell {
-	const cellMetadata = cell.metadata?.custom as CellMetadata | undefined;
+	const cellMetadata = getCellMetadata(cell);
 	const codeCell: any = {
 		cell_type: 'code',
 		execution_count: cell.executionSummary?.executionOrder ?? null,
@@ -69,7 +72,7 @@ function createCodeCellFromNotebookCell(cell: NotebookCellData): nbformat.ICodeC
 }
 
 function createRawCellFromNotebookCell(cell: NotebookCellData): nbformat.IRawCell {
-	const cellMetadata = cell.metadata?.custom as CellMetadata | undefined;
+	const cellMetadata = getCellMetadata(cell);
 	const rawCell: any = {
 		cell_type: 'raw',
 		source: splitMultilineString(cell.value.replace(/\r\n/g, '\n')),
@@ -319,7 +322,7 @@ function convertOutputMimeToJupyterOutput(mime: string, value: Uint8Array) {
 }
 
 function createMarkdownCellFromNotebookCell(cell: NotebookCellData): nbformat.IMarkdownCell {
-	const cellMetadata = cell.metadata?.custom as CellMetadata | undefined;
+	const cellMetadata = getCellMetadata(cell);
 	const markdownCell: any = {
 		cell_type: 'markdown',
 		source: splitMultilineString(cell.value.replace(/\r\n/g, '\n')),

extensions/ipynb/yarn.lock

@@ -76,6 +76,11 @@
   dependencies:
     "@phosphor/algorithm" "^1.2.0"
 
+"@types/uuid@^8.3.1":
+  version "8.3.1"
+  resolved "https://registry.yarnpkg.com/@types/uuid/-/uuid-8.3.1.tgz#1a32969cf8f0364b3d8c8af9cc3555b7805df14f"
+  integrity sha512-Y2mHTRAbqfFkpjldbkHGY8JIzRN6XqYRliG8/24FcHm2D2PwW24fl5xMRTVGdrb7iMrwCaIEbLWerGIkXuFWVg==
+
 ajv@^6.5.5:
   version "6.12.6"
   resolved "https://registry.yarnpkg.com/ajv/-/ajv-6.12.6.tgz#baf5a62e802b07d977034586f8c3baf5adf26df4"
@@ -157,3 +162,8 @@ url-parse@~1.4.3:
   dependencies:
     querystringify "^2.1.1"
     requires-port "^1.0.0"
+
+uuid@^8.3.2:
+  version "8.3.2"
+  resolved "https://registry.yarnpkg.com/uuid/-/uuid-8.3.2.tgz#80d5b5ced271bb9af6c445f21a1a04c606cefbe2"
+  integrity sha512-+NYs2QeMWy+GWFOEm9xnn6HCDp0l7QBD7ml8zLUmJ+93Q5NF0NocErnwkTkXVFNiX3/fpC6afS8Dhb/gz7R7eg==