Skip to content

Commit e3cf253

Browse files
DonJayamanneDonJayamanne
authored
Include execution count in jupyter notebook diff (microsoft#208292)
* Include execution count in notebook diff * Fix tests * Misc changes * Fix tests * Fix more tests
1 parent 5e0394c commit e3cf253

File tree

6 files changed

+111
-46
lines changed

6 files changed

+111
-46
lines changed

‎extensions/ipynb/src/common.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,10 @@ export interface CellMetadata {
6060
* Stores cell metadata.
6161
*/
6262
metadata?: Partial<nbformat.ICellMetadata> & { vscode?: { languageId?: string } };
63+
/**
64+
* The code cell's prompt number. Will be null if the cell has not been run.
65+
*/
66+
execution_count?: number;
6367
}
6468

6569
export function useCustomPropertyInMetadata() {

‎extensions/ipynb/src/deserializers.ts

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,11 @@ function getNotebookCellMetadata(cell: nbformat.IBaseCell): {
159159
// We put this only for VSC to display in diff view.
160160
// Else we don't use this.
161161
const custom: CellMetadata = {};
162+
163+
if (cell.cell_type === 'code' && typeof cell['execution_count'] === 'number') {
164+
custom.execution_count = cell['execution_count'];
165+
}
166+
162167
if (cell['metadata']) {
163168
custom['metadata'] = JSON.parse(JSON.stringify(cell['metadata']));
164169
}
@@ -177,6 +182,10 @@ function getNotebookCellMetadata(cell: nbformat.IBaseCell): {
177182
// We put this only for VSC to display in diff view.
178183
// Else we don't use this.
179184
const cellMetadata: CellMetadata = {};
185+
if (cell.cell_type === 'code' && typeof cell['execution_count'] === 'number') {
186+
cellMetadata.execution_count = cell['execution_count'];
187+
}
188+
180189
if (cell['metadata']) {
181190
cellMetadata['metadata'] = JSON.parse(JSON.stringify(cell['metadata']));
182191
}

‎extensions/ipynb/src/notebookModelStoreSync.ts

Lines changed: 49 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*--------------------------------------------------------------------------------------------*/
55

66
import { ExtensionContext, NotebookCellKind, NotebookDocument, NotebookDocumentChangeEvent, NotebookEdit, workspace, WorkspaceEdit, type NotebookCell, type NotebookDocumentWillSaveEvent } from 'vscode';
7-
import { getCellMetadata, getVSCodeCellLanguageId, removeVSCodeCellLanguageId, setVSCodeCellLanguageId } from './serializers';
7+
import { getCellMetadata, getVSCodeCellLanguageId, removeVSCodeCellLanguageId, setVSCodeCellLanguageId, sortObjectPropertiesRecursively } from './serializers';
88
import { CellMetadata, useCustomPropertyInMetadata } from './common';
99
import { getNotebookMetadata } from './notebookSerializer';
1010
import type * as nbformat from '@jupyterlab/nbformat';
@@ -53,15 +53,25 @@ function cleanup(notebook: NotebookDocument, promise: PromiseLike<void>) {
5353
}
5454
}
5555
}
56-
function trackAndUpdateCellMetadata(notebook: NotebookDocument, cell: NotebookCell, metadata: CellMetadata & { vscode?: { languageId: string } }) {
56+
function trackAndUpdateCellMetadata(notebook: NotebookDocument, updates: { cell: NotebookCell; metadata: CellMetadata & { vscode?: { languageId: string } } }[]) {
5757
const pendingUpdates = pendingNotebookCellModelUpdates.get(notebook) ?? new Set<Thenable<void>>();
5858
pendingNotebookCellModelUpdates.set(notebook, pendingUpdates);
5959
const edit = new WorkspaceEdit();
60-
if (useCustomPropertyInMetadata()) {
61-
edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, { ...(cell.metadata), custom: metadata })]);
62-
} else {
63-
edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, { ...cell.metadata, ...metadata })]);
64-
}
60+
updates.forEach(({ cell, metadata }) => {
61+
let newMetadata: any = {};
62+
if (useCustomPropertyInMetadata()) {
63+
newMetadata = { ...(cell.metadata), custom: metadata };
64+
} else {
65+
newMetadata = { ...cell.metadata, ...metadata };
66+
if (!metadata.execution_count && newMetadata.execution_count) {
67+
delete newMetadata.execution_count;
68+
}
69+
if (!metadata.attachments && newMetadata.attachments) {
70+
delete newMetadata.attachments;
71+
}
72+
}
73+
edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, sortObjectPropertiesRecursively(newMetadata))]);
74+
});
6575
const promise = workspace.applyEdit(edit).then(noop, noop);
6676
pendingUpdates.add(promise);
6777
const clean = () => cleanup(notebook, promise);
@@ -78,31 +88,41 @@ function onDidChangeNotebookCells(e: NotebookDocumentChangeEvent) {
7888

7989
// use the preferred language from document metadata or the first cell language as the notebook preferred cell language
8090
const preferredCellLanguage = notebookMetadata.metadata?.language_info?.name;
81-
91+
const updates: { cell: NotebookCell; metadata: CellMetadata & { vscode?: { languageId: string } } }[] = [];
8292
// When we change the language of a cell,
8393
// Ensure the metadata in the notebook cell has been updated as well,
8494
// Else model will be out of sync with ipynb https://github.com/microsoft/vscode/issues/207968#issuecomment-2002858596
8595
e.cellChanges.forEach(e => {
8696
if (!preferredCellLanguage || e.cell.kind !== NotebookCellKind.Code) {
8797
return;
8898
}
89-
const languageIdInMetadata = getVSCodeCellLanguageId(getCellMetadata(e.cell));
90-
if (e.cell.document.languageId !== preferredCellLanguage && e.cell.document.languageId !== languageIdInMetadata) {
91-
const metadata: CellMetadata = JSON.parse(JSON.stringify(getCellMetadata(e.cell)));
92-
metadata.metadata = metadata.metadata || {};
93-
setVSCodeCellLanguageId(metadata, e.cell.document.languageId);
94-
trackAndUpdateCellMetadata(notebook, e.cell, metadata);
99+
const currentMetadata = e.metadata ? getCellMetadata({ metadata: e.metadata }) : getCellMetadata({ cell: e.cell });
100+
const languageIdInMetadata = getVSCodeCellLanguageId(currentMetadata);
101+
const metadata: CellMetadata = JSON.parse(JSON.stringify(currentMetadata));
102+
metadata.metadata = metadata.metadata || {};
103+
let metadataUpdated = false;
104+
if (e.executionSummary?.executionOrder && typeof e.executionSummary.success === 'boolean' && currentMetadata.execution_count !== e.executionSummary?.executionOrder) {
105+
metadata.execution_count = e.executionSummary.executionOrder;
106+
metadataUpdated = true;
107+
} else if (!e.executionSummary && !e.metadata && e.outputs?.length === 0 && currentMetadata.execution_count) {
108+
// Clear all.
109+
delete metadata.execution_count;
110+
metadataUpdated = true;
111+
}
95112

96-
} else if (e.cell.document.languageId === preferredCellLanguage && languageIdInMetadata) {
97-
const metadata: CellMetadata = JSON.parse(JSON.stringify(getCellMetadata(e.cell)));
98-
metadata.metadata = metadata.metadata || {};
113+
if (e.document?.languageId && e.document?.languageId !== preferredCellLanguage && e.document?.languageId !== languageIdInMetadata) {
114+
setVSCodeCellLanguageId(metadata, e.document.languageId);
115+
metadataUpdated = true;
116+
} else if (e.document?.languageId && e.document.languageId === preferredCellLanguage && languageIdInMetadata) {
99117
removeVSCodeCellLanguageId(metadata);
100-
trackAndUpdateCellMetadata(notebook, e.cell, metadata);
101-
} else if (e.cell.document.languageId === preferredCellLanguage && e.cell.document.languageId === languageIdInMetadata) {
102-
const metadata: CellMetadata = JSON.parse(JSON.stringify(getCellMetadata(e.cell)));
103-
metadata.metadata = metadata.metadata || {};
118+
metadataUpdated = true;
119+
} else if (e.document?.languageId && e.document.languageId === preferredCellLanguage && e.document.languageId === languageIdInMetadata) {
104120
removeVSCodeCellLanguageId(metadata);
105-
trackAndUpdateCellMetadata(notebook, e.cell, metadata);
121+
metadataUpdated = true;
122+
}
123+
124+
if (metadataUpdated) {
125+
updates.push({ cell: e.cell, metadata });
106126
}
107127
});
108128

@@ -112,7 +132,7 @@ function onDidChangeNotebookCells(e: NotebookDocumentChangeEvent) {
112132
change.addedCells.forEach(cell => {
113133
// When ever a cell is added, always update the metadata
114134
// as metadata is always an empty `{}` in ipynb JSON file
115-
const cellMetadata = getCellMetadata(cell);
135+
const cellMetadata = getCellMetadata({ cell });
116136

117137
// Avoid updating the metadata if it's not required.
118138
if (cellMetadata.metadata) {
@@ -131,9 +151,13 @@ function onDidChangeNotebookCells(e: NotebookDocumentChangeEvent) {
131151
if (isCellIdRequired(notebookMetadata) && !cellMetadata?.id) {
132152
metadata.id = generateCellId(e.notebook);
133153
}
134-
trackAndUpdateCellMetadata(notebook, cell, metadata);
154+
updates.push({ cell, metadata });
135155
});
136156
});
157+
158+
if (updates.length) {
159+
trackAndUpdateCellMetadata(notebook, updates);
160+
}
137161
}
138162

139163

@@ -158,7 +182,7 @@ function generateCellId(notebook: NotebookDocument) {
158182
let duplicate = false;
159183
for (let index = 0; index < notebook.cellCount; index++) {
160184
const cell = notebook.cellAt(index);
161-
const existingId = getCellMetadata(cell)?.id;
185+
const existingId = getCellMetadata({ cell })?.id;
162186
if (!existingId) {
163187
continue;
164188
}

‎extensions/ipynb/src/serializers.ts

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -54,26 +54,48 @@ export function sortObjectPropertiesRecursively(obj: any): any {
5454
return obj;
5555
}
5656

57-
export function getCellMetadata(cell: NotebookCell | NotebookCellData): CellMetadata {
58-
if (useCustomPropertyInMetadata()) {
57+
export function getCellMetadata(options: { cell: NotebookCell | NotebookCellData } | { metadata?: { [key: string]: any } }): CellMetadata {
58+
if ('cell' in options) {
59+
const cell = options.cell;
60+
if (useCustomPropertyInMetadata()) {
61+
const metadata: CellMetadata = {
62+
// it contains the cell id, and the cell metadata, along with other nb cell metadata
63+
...(cell.metadata?.custom ?? {})
64+
};
65+
// promote the cell attachments to the top level
66+
const attachments = cell.metadata?.custom?.attachments ?? cell.metadata?.attachments;
67+
if (attachments) {
68+
metadata.attachments = attachments;
69+
}
70+
return metadata;
71+
}
5972
const metadata = {
6073
// it contains the cell id, and the cell metadata, along with other nb cell metadata
61-
...(cell.metadata?.custom ?? {})
74+
...(cell.metadata ?? {})
6275
};
6376

64-
// promote the cell attachments to the top level
65-
const attachments = cell.metadata?.custom?.attachments ?? cell.metadata?.attachments;
66-
if (attachments) {
67-
metadata.attachments = attachments;
77+
return metadata;
78+
} else {
79+
const cell = options;
80+
if (useCustomPropertyInMetadata()) {
81+
const metadata: CellMetadata = {
82+
// it contains the cell id, and the cell metadata, along with other nb cell metadata
83+
...(cell.metadata?.custom ?? {})
84+
};
85+
// promote the cell attachments to the top level
86+
const attachments = cell.metadata?.custom?.attachments ?? cell.metadata?.attachments;
87+
if (attachments) {
88+
metadata.attachments = attachments;
89+
}
90+
return metadata;
6891
}
92+
const metadata = {
93+
// it contains the cell id, and the cell metadata, along with other nb cell metadata
94+
...(cell.metadata ?? {})
95+
};
96+
6997
return metadata;
7098
}
71-
const metadata = {
72-
// it contains the cell id, and the cell metadata, along with other nb cell metadata
73-
...(cell.metadata ?? {})
74-
};
75-
76-
return metadata;
7799
}
78100

79101
export function getVSCodeCellLanguageId(metadata: CellMetadata): string | undefined {
@@ -90,7 +112,7 @@ export function removeVSCodeCellLanguageId(metadata: CellMetadata) {
90112
}
91113

92114
function createCodeCellFromNotebookCell(cell: NotebookCellData, preferredLanguage: string | undefined): nbformat.ICodeCell {
93-
const cellMetadata: CellMetadata = JSON.parse(JSON.stringify(getCellMetadata(cell)));
115+
const cellMetadata: CellMetadata = JSON.parse(JSON.stringify(getCellMetadata({ cell })));
94116
cellMetadata.metadata = cellMetadata.metadata || {}; // This cannot be empty.
95117
if (cell.languageId !== preferredLanguage) {
96118
setVSCodeCellLanguageId(cellMetadata, cell.languageId);
@@ -113,7 +135,7 @@ function createCodeCellFromNotebookCell(cell: NotebookCellData, preferredLanguag
113135
}
114136

115137
function createRawCellFromNotebookCell(cell: NotebookCellData): nbformat.IRawCell {
116-
const cellMetadata = getCellMetadata(cell);
138+
const cellMetadata = getCellMetadata({ cell });
117139
const rawCell: any = {
118140
cell_type: 'raw',
119141
source: splitMultilineString(cell.value.replace(/\r\n/g, '\n')),
@@ -364,7 +386,7 @@ function convertOutputMimeToJupyterOutput(mime: string, value: Uint8Array) {
364386
}
365387

366388
export function createMarkdownCellFromNotebookCell(cell: NotebookCellData): nbformat.IMarkdownCell {
367-
const cellMetadata = getCellMetadata(cell);
389+
const cellMetadata = getCellMetadata({ cell });
368390
const markdownCell: any = {
369391
cell_type: 'markdown',
370392
source: splitMultilineString(cell.value.replace(/\r\n/g, '\n')),

‎extensions/ipynb/src/test/notebookModelStoreSync.test.ts

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -330,7 +330,9 @@ import { activate } from '../notebookModelStoreSync';
330330
cellChanges: [
331331
{
332332
cell,
333-
document: undefined,
333+
document: {
334+
languageId: 'javascript'
335+
} as any,
334336
metadata: undefined,
335337
outputs: undefined,
336338
executionSummary: undefined
@@ -465,7 +467,9 @@ import { activate } from '../notebookModelStoreSync';
465467
cellChanges: [
466468
{
467469
cell,
468-
document: undefined,
470+
document: {
471+
languageId: 'javascript'
472+
} as any,
469473
metadata: undefined,
470474
outputs: undefined,
471475
executionSummary: undefined
@@ -540,7 +544,9 @@ import { activate } from '../notebookModelStoreSync';
540544
cellChanges: [
541545
{
542546
cell,
543-
document: undefined,
547+
document: {
548+
languageId: 'powershell'
549+
} as any,
544550
metadata: undefined,
545551
outputs: undefined,
546552
executionSummary: undefined

‎extensions/ipynb/src/test/serializers.test.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ function deepStripProperties(obj: any, props: string[]) {
6464

6565
const expectedCodeCell = new vscode.NotebookCellData(vscode.NotebookCellKind.Code, 'print(1)', 'python');
6666
expectedCodeCell.outputs = [];
67-
expectedCodeCell.metadata = useCustomPropertyInMetadata ? { custom: { metadata: {} } } : { metadata: {} };
67+
expectedCodeCell.metadata = useCustomPropertyInMetadata ? { custom: { execution_count: 10, metadata: {} } } : { execution_count: 10, metadata: {} };
6868
expectedCodeCell.executionSummary = { executionOrder: 10 };
6969

7070
const expectedMarkdownCell = new vscode.NotebookCellData(vscode.NotebookCellKind.Markup, '# HEAD', 'markdown');
@@ -105,7 +105,7 @@ function deepStripProperties(obj: any, props: string[]) {
105105
}
106106
};
107107

108-
const cellMetadata = getCellMetadata(markdownCell);
108+
const cellMetadata = getCellMetadata({ cell: markdownCell });
109109
assert.deepStrictEqual(cellMetadata, {
110110
id: '123',
111111
metadata: {

0 commit comments

Comments
 (0)