Include execution count in jupyter notebook diff (#208292)

* Include execution count in notebook diff

* Fix tests

* Misc changes

* Fix tests

* Fix more tests
This commit is contained in:
Don Jayamanne
2024-03-22 00:14:33 +11:00
committed by GitHub
parent 5e0394c8d1
commit e3cf2530e6
6 changed files with 111 additions and 46 deletions

View File

@@ -60,6 +60,10 @@ export interface CellMetadata {
* Stores cell metadata.
*/
metadata?: Partial<nbformat.ICellMetadata> & { vscode?: { languageId?: string } };
/**
* The code cell's prompt number. Will be null if the cell has not been run.
*/
execution_count?: number;
}
export function useCustomPropertyInMetadata() {

View File

@@ -159,6 +159,11 @@ function getNotebookCellMetadata(cell: nbformat.IBaseCell): {
// We put this only for VSC to display in diff view.
// Else we don't use this.
const custom: CellMetadata = {};
if (cell.cell_type === 'code' && typeof cell['execution_count'] === 'number') {
custom.execution_count = cell['execution_count'];
}
if (cell['metadata']) {
custom['metadata'] = JSON.parse(JSON.stringify(cell['metadata']));
}
@@ -177,6 +182,10 @@ function getNotebookCellMetadata(cell: nbformat.IBaseCell): {
// We put this only for VSC to display in diff view.
// Else we don't use this.
const cellMetadata: CellMetadata = {};
if (cell.cell_type === 'code' && typeof cell['execution_count'] === 'number') {
cellMetadata.execution_count = cell['execution_count'];
}
if (cell['metadata']) {
cellMetadata['metadata'] = JSON.parse(JSON.stringify(cell['metadata']));
}

View File

@@ -4,7 +4,7 @@
*--------------------------------------------------------------------------------------------*/
import { ExtensionContext, NotebookCellKind, NotebookDocument, NotebookDocumentChangeEvent, NotebookEdit, workspace, WorkspaceEdit, type NotebookCell, type NotebookDocumentWillSaveEvent } from 'vscode';
import { getCellMetadata, getVSCodeCellLanguageId, removeVSCodeCellLanguageId, setVSCodeCellLanguageId } from './serializers';
import { getCellMetadata, getVSCodeCellLanguageId, removeVSCodeCellLanguageId, setVSCodeCellLanguageId, sortObjectPropertiesRecursively } from './serializers';
import { CellMetadata, useCustomPropertyInMetadata } from './common';
import { getNotebookMetadata } from './notebookSerializer';
import type * as nbformat from '@jupyterlab/nbformat';
@@ -53,15 +53,25 @@ function cleanup(notebook: NotebookDocument, promise: PromiseLike<void>) {
}
}
}
function trackAndUpdateCellMetadata(notebook: NotebookDocument, cell: NotebookCell, metadata: CellMetadata & { vscode?: { languageId: string } }) {
function trackAndUpdateCellMetadata(notebook: NotebookDocument, updates: { cell: NotebookCell; metadata: CellMetadata & { vscode?: { languageId: string } } }[]) {
const pendingUpdates = pendingNotebookCellModelUpdates.get(notebook) ?? new Set<Thenable<void>>();
pendingNotebookCellModelUpdates.set(notebook, pendingUpdates);
const edit = new WorkspaceEdit();
if (useCustomPropertyInMetadata()) {
edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, { ...(cell.metadata), custom: metadata })]);
} else {
edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, { ...cell.metadata, ...metadata })]);
}
updates.forEach(({ cell, metadata }) => {
let newMetadata: any = {};
if (useCustomPropertyInMetadata()) {
newMetadata = { ...(cell.metadata), custom: metadata };
} else {
newMetadata = { ...cell.metadata, ...metadata };
if (!metadata.execution_count && newMetadata.execution_count) {
delete newMetadata.execution_count;
}
if (!metadata.attachments && newMetadata.attachments) {
delete newMetadata.attachments;
}
}
edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, sortObjectPropertiesRecursively(newMetadata))]);
});
const promise = workspace.applyEdit(edit).then(noop, noop);
pendingUpdates.add(promise);
const clean = () => cleanup(notebook, promise);
@@ -78,7 +88,7 @@ function onDidChangeNotebookCells(e: NotebookDocumentChangeEvent) {
// use the preferred language from document metadata or the first cell language as the notebook preferred cell language
const preferredCellLanguage = notebookMetadata.metadata?.language_info?.name;
const updates: { cell: NotebookCell; metadata: CellMetadata & { vscode?: { languageId: string } } }[] = [];
// When we change the language of a cell,
// Ensure the metadata in the notebook cell has been updated as well,
// Else model will be out of sync with ipynb https://github.com/microsoft/vscode/issues/207968#issuecomment-2002858596
@@ -86,23 +96,33 @@ function onDidChangeNotebookCells(e: NotebookDocumentChangeEvent) {
if (!preferredCellLanguage || e.cell.kind !== NotebookCellKind.Code) {
return;
}
const languageIdInMetadata = getVSCodeCellLanguageId(getCellMetadata(e.cell));
if (e.cell.document.languageId !== preferredCellLanguage && e.cell.document.languageId !== languageIdInMetadata) {
const metadata: CellMetadata = JSON.parse(JSON.stringify(getCellMetadata(e.cell)));
metadata.metadata = metadata.metadata || {};
setVSCodeCellLanguageId(metadata, e.cell.document.languageId);
trackAndUpdateCellMetadata(notebook, e.cell, metadata);
const currentMetadata = e.metadata ? getCellMetadata({ metadata: e.metadata }) : getCellMetadata({ cell: e.cell });
const languageIdInMetadata = getVSCodeCellLanguageId(currentMetadata);
const metadata: CellMetadata = JSON.parse(JSON.stringify(currentMetadata));
metadata.metadata = metadata.metadata || {};
let metadataUpdated = false;
if (e.executionSummary?.executionOrder && typeof e.executionSummary.success === 'boolean' && currentMetadata.execution_count !== e.executionSummary?.executionOrder) {
metadata.execution_count = e.executionSummary.executionOrder;
metadataUpdated = true;
} else if (!e.executionSummary && !e.metadata && e.outputs?.length === 0 && currentMetadata.execution_count) {
// Clear all.
delete metadata.execution_count;
metadataUpdated = true;
}
} else if (e.cell.document.languageId === preferredCellLanguage && languageIdInMetadata) {
const metadata: CellMetadata = JSON.parse(JSON.stringify(getCellMetadata(e.cell)));
metadata.metadata = metadata.metadata || {};
if (e.document?.languageId && e.document?.languageId !== preferredCellLanguage && e.document?.languageId !== languageIdInMetadata) {
setVSCodeCellLanguageId(metadata, e.document.languageId);
metadataUpdated = true;
} else if (e.document?.languageId && e.document.languageId === preferredCellLanguage && languageIdInMetadata) {
removeVSCodeCellLanguageId(metadata);
trackAndUpdateCellMetadata(notebook, e.cell, metadata);
} else if (e.cell.document.languageId === preferredCellLanguage && e.cell.document.languageId === languageIdInMetadata) {
const metadata: CellMetadata = JSON.parse(JSON.stringify(getCellMetadata(e.cell)));
metadata.metadata = metadata.metadata || {};
metadataUpdated = true;
} else if (e.document?.languageId && e.document.languageId === preferredCellLanguage && e.document.languageId === languageIdInMetadata) {
removeVSCodeCellLanguageId(metadata);
trackAndUpdateCellMetadata(notebook, e.cell, metadata);
metadataUpdated = true;
}
if (metadataUpdated) {
updates.push({ cell: e.cell, metadata });
}
});
@@ -112,7 +132,7 @@ function onDidChangeNotebookCells(e: NotebookDocumentChangeEvent) {
change.addedCells.forEach(cell => {
// When ever a cell is added, always update the metadata
// as metadata is always an empty `{}` in ipynb JSON file
const cellMetadata = getCellMetadata(cell);
const cellMetadata = getCellMetadata({ cell });
// Avoid updating the metadata if it's not required.
if (cellMetadata.metadata) {
@@ -131,9 +151,13 @@ function onDidChangeNotebookCells(e: NotebookDocumentChangeEvent) {
if (isCellIdRequired(notebookMetadata) && !cellMetadata?.id) {
metadata.id = generateCellId(e.notebook);
}
trackAndUpdateCellMetadata(notebook, cell, metadata);
updates.push({ cell, metadata });
});
});
if (updates.length) {
trackAndUpdateCellMetadata(notebook, updates);
}
}
@@ -158,7 +182,7 @@ function generateCellId(notebook: NotebookDocument) {
let duplicate = false;
for (let index = 0; index < notebook.cellCount; index++) {
const cell = notebook.cellAt(index);
const existingId = getCellMetadata(cell)?.id;
const existingId = getCellMetadata({ cell })?.id;
if (!existingId) {
continue;
}

View File

@@ -54,26 +54,48 @@ export function sortObjectPropertiesRecursively(obj: any): any {
return obj;
}
export function getCellMetadata(cell: NotebookCell | NotebookCellData): CellMetadata {
if (useCustomPropertyInMetadata()) {
export function getCellMetadata(options: { cell: NotebookCell | NotebookCellData } | { metadata?: { [key: string]: any } }): CellMetadata {
if ('cell' in options) {
const cell = options.cell;
if (useCustomPropertyInMetadata()) {
const metadata: CellMetadata = {
// it contains the cell id, and the cell metadata, along with other nb cell metadata
...(cell.metadata?.custom ?? {})
};
// promote the cell attachments to the top level
const attachments = cell.metadata?.custom?.attachments ?? cell.metadata?.attachments;
if (attachments) {
metadata.attachments = attachments;
}
return metadata;
}
const metadata = {
// it contains the cell id, and the cell metadata, along with other nb cell metadata
...(cell.metadata?.custom ?? {})
...(cell.metadata ?? {})
};
// promote the cell attachments to the top level
const attachments = cell.metadata?.custom?.attachments ?? cell.metadata?.attachments;
if (attachments) {
metadata.attachments = attachments;
return metadata;
} else {
const cell = options;
if (useCustomPropertyInMetadata()) {
const metadata: CellMetadata = {
// it contains the cell id, and the cell metadata, along with other nb cell metadata
...(cell.metadata?.custom ?? {})
};
// promote the cell attachments to the top level
const attachments = cell.metadata?.custom?.attachments ?? cell.metadata?.attachments;
if (attachments) {
metadata.attachments = attachments;
}
return metadata;
}
const metadata = {
// it contains the cell id, and the cell metadata, along with other nb cell metadata
...(cell.metadata ?? {})
};
return metadata;
}
const metadata = {
// it contains the cell id, and the cell metadata, along with other nb cell metadata
...(cell.metadata ?? {})
};
return metadata;
}
export function getVSCodeCellLanguageId(metadata: CellMetadata): string | undefined {
@@ -90,7 +112,7 @@ export function removeVSCodeCellLanguageId(metadata: CellMetadata) {
}
function createCodeCellFromNotebookCell(cell: NotebookCellData, preferredLanguage: string | undefined): nbformat.ICodeCell {
const cellMetadata: CellMetadata = JSON.parse(JSON.stringify(getCellMetadata(cell)));
const cellMetadata: CellMetadata = JSON.parse(JSON.stringify(getCellMetadata({ cell })));
cellMetadata.metadata = cellMetadata.metadata || {}; // This cannot be empty.
if (cell.languageId !== preferredLanguage) {
setVSCodeCellLanguageId(cellMetadata, cell.languageId);
@@ -113,7 +135,7 @@ function createCodeCellFromNotebookCell(cell: NotebookCellData, preferredLanguag
}
function createRawCellFromNotebookCell(cell: NotebookCellData): nbformat.IRawCell {
const cellMetadata = getCellMetadata(cell);
const cellMetadata = getCellMetadata({ cell });
const rawCell: any = {
cell_type: 'raw',
source: splitMultilineString(cell.value.replace(/\r\n/g, '\n')),
@@ -364,7 +386,7 @@ function convertOutputMimeToJupyterOutput(mime: string, value: Uint8Array) {
}
export function createMarkdownCellFromNotebookCell(cell: NotebookCellData): nbformat.IMarkdownCell {
const cellMetadata = getCellMetadata(cell);
const cellMetadata = getCellMetadata({ cell });
const markdownCell: any = {
cell_type: 'markdown',
source: splitMultilineString(cell.value.replace(/\r\n/g, '\n')),

View File

@@ -330,7 +330,9 @@ import { activate } from '../notebookModelStoreSync';
cellChanges: [
{
cell,
document: undefined,
document: {
languageId: 'javascript'
} as any,
metadata: undefined,
outputs: undefined,
executionSummary: undefined
@@ -465,7 +467,9 @@ import { activate } from '../notebookModelStoreSync';
cellChanges: [
{
cell,
document: undefined,
document: {
languageId: 'javascript'
} as any,
metadata: undefined,
outputs: undefined,
executionSummary: undefined
@@ -540,7 +544,9 @@ import { activate } from '../notebookModelStoreSync';
cellChanges: [
{
cell,
document: undefined,
document: {
languageId: 'powershell'
} as any,
metadata: undefined,
outputs: undefined,
executionSummary: undefined

View File

@@ -64,7 +64,7 @@ function deepStripProperties(obj: any, props: string[]) {
const expectedCodeCell = new vscode.NotebookCellData(vscode.NotebookCellKind.Code, 'print(1)', 'python');
expectedCodeCell.outputs = [];
expectedCodeCell.metadata = useCustomPropertyInMetadata ? { custom: { metadata: {} } } : { metadata: {} };
expectedCodeCell.metadata = useCustomPropertyInMetadata ? { custom: { execution_count: 10, metadata: {} } } : { execution_count: 10, metadata: {} };
expectedCodeCell.executionSummary = { executionOrder: 10 };
const expectedMarkdownCell = new vscode.NotebookCellData(vscode.NotebookCellKind.Markup, '# HEAD', 'markdown');
@@ -105,7 +105,7 @@ function deepStripProperties(obj: any, props: string[]) {
}
};
const cellMetadata = getCellMetadata(markdownCell);
const cellMetadata = getCellMetadata({ cell: markdownCell });
assert.deepStrictEqual(cellMetadata, {
id: '123',
metadata: {