Skip to content

Commit 6e8bc02

Browse files
rebornixrebornix
authored
Cell attachment cleanup tool improvement in diff editor (microsoft#161132)
* Move attachment out of custom metadata, prep for attachment clean up in diff editor * recover attachments from dirty notebook document * Allow metadata to be restored when content changed/reverted in nb diff editor
1 parent 2d7655c commit 6e8bc02

19 files changed

+166
-47
lines changed

‎extensions/ipynb/notebook-src/cellAttachmentRenderer.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ export async function activate(ctx: RendererContext<void>) {
2222
md.renderer.rules.image = (tokens: MarkdownItToken[], idx: number, options, env, self) => {
2323
const token = tokens[idx];
2424
const src = token.attrGet('src');
25-
const attachments: Record<string, Record<string, string>> = env.outputItem.metadata.custom?.attachments;
25+
const attachments: Record<string, Record<string, string>> = env.outputItem.metadata.attachments;
2626
if (attachments && src) {
2727
const imageAttachment = attachments[src.replace('attachment:', '')];
2828
if (imageAttachment) {

‎extensions/ipynb/package.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,8 @@
99
"vscode": "^1.57.0"
1010
},
1111
"enabledApiProposals": [
12-
"documentPaste"
12+
"documentPaste",
13+
"diffContentOptions"
1314
],
1415
"activationEvents": [
1516
"*"

‎extensions/ipynb/src/deserializers.ts

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -149,21 +149,29 @@ function convertJupyterOutputToBuffer(mime: string, value: unknown): NotebookCel
149149
}
150150
}
151151

152-
function getNotebookCellMetadata(cell: nbformat.IBaseCell): CellMetadata {
152+
function getNotebookCellMetadata(cell: nbformat.IBaseCell): {
153+
[key: string]: any;
154+
} {
155+
const cellMetadata: { [key: string]: any } = {};
153156
// We put this only for VSC to display in diff view.
154157
// Else we don't use this.
155-
const propertiesToClone: (keyof CellMetadata)[] = ['metadata', 'attachments'];
156158
const custom: CellMetadata = {};
157-
propertiesToClone.forEach((propertyToClone) => {
158-
if (cell[propertyToClone]) {
159-
custom[propertyToClone] = JSON.parse(JSON.stringify(cell[propertyToClone]));
160-
}
161-
});
159+
if (cell['metadata']) {
160+
custom['metadata'] = JSON.parse(JSON.stringify(cell['metadata']));
161+
}
162+
162163
if ('id' in cell && typeof cell.id === 'string') {
163164
custom.id = cell.id;
164165
}
165-
return custom;
166+
167+
cellMetadata.custom = custom;
168+
169+
if (cell['attachments']) {
170+
cellMetadata.attachments = JSON.parse(JSON.stringify(cell['attachments']));
171+
}
172+
return cellMetadata;
166173
}
174+
167175
function getOutputMetadata(output: nbformat.IOutput): CellOutputMetadata {
168176
// Add on transient data if we have any. This should be removed by our save functions elsewhere.
169177
const metadata: CellOutputMetadata = {
@@ -284,7 +292,7 @@ export function jupyterCellOutputToCellOutput(output: nbformat.IOutput): Noteboo
284292
function createNotebookCellDataFromRawCell(cell: nbformat.IRawCell): NotebookCellData {
285293
const cellData = new NotebookCellData(NotebookCellKind.Code, concatMultilineString(cell.source), 'raw');
286294
cellData.outputs = [];
287-
cellData.metadata = { custom: getNotebookCellMetadata(cell) };
295+
cellData.metadata = getNotebookCellMetadata(cell);
288296
return cellData;
289297
}
290298
function createNotebookCellDataFromMarkdownCell(cell: nbformat.IMarkdownCell): NotebookCellData {
@@ -294,7 +302,7 @@ function createNotebookCellDataFromMarkdownCell(cell: nbformat.IMarkdownCell): N
294302
'markdown'
295303
);
296304
cellData.outputs = [];
297-
cellData.metadata = { custom: getNotebookCellMetadata(cell) };
305+
cellData.metadata = getNotebookCellMetadata(cell);
298306
return cellData;
299307
}
300308
function createNotebookCellDataFromCodeCell(cell: nbformat.ICodeCell, cellLanguage: string): NotebookCellData {
@@ -313,7 +321,7 @@ function createNotebookCellDataFromCodeCell(cell: nbformat.ICodeCell, cellLangua
313321
const cellData = new NotebookCellData(NotebookCellKind.Code, source, cellLanguageId);
314322

315323
cellData.outputs = outputs;
316-
cellData.metadata = { custom: getNotebookCellMetadata(cell) };
324+
cellData.metadata = getNotebookCellMetadata(cell);
317325
cellData.executionSummary = executionSummary;
318326
return cellData;
319327
}

‎extensions/ipynb/src/ipynbMain.ts

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,13 @@ export function activate(context: vscode.ExtensionContext) {
3535
transientOutputs: false,
3636
transientCellMetadata: {
3737
breakpointMargin: true,
38-
custom: false
38+
custom: false,
39+
attachments: false
40+
},
41+
cellContentMetadata: {
42+
attachments: true
3943
}
40-
}));
44+
} as vscode.NotebookDocumentContentOptions));
4145

4246
vscode.languages.registerCodeLensProvider({ pattern: '**/*.ipynb' }, {
4347
provideCodeLenses: (document) => {

‎extensions/ipynb/src/notebookAttachmentCleaner.ts

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -156,13 +156,13 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
156156
if (this.checkMetadataAttachmentsExistence(cell.metadata)) {
157157
// the cell metadata contains attachments, check if any are used in the markdown source
158158

159-
for (const currFilename of Object.keys(cell.metadata.custom.attachments)) {
159+
for (const currFilename of Object.keys(cell.metadata.attachments)) {
160160
// means markdown reference is present in the metadata, rendering will work properly
161161
// therefore, we don't need to check it in the next loop either
162162
if (markdownAttachmentsRefedInCell.has(currFilename)) {
163163
// attachment reference is present in the markdown source, no need to cache it
164164
markdownAttachmentsRefedInCell.get(currFilename)!.valid = true;
165-
markdownAttachmentsInUse[currFilename] = cell.metadata.custom.attachments[currFilename];
165+
markdownAttachmentsInUse[currFilename] = cell.metadata.attachments[currFilename];
166166
} else {
167167
// attachment reference is not present in the markdown source, cache it
168168
this.saveAttachmentToCache(notebookUri, cellFragment, currFilename, cell.metadata);
@@ -187,9 +187,9 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
187187
}
188188
}
189189

190-
if (!objectEquals(markdownAttachmentsInUse, cell.metadata.custom.attachments)) {
190+
if (!objectEquals(markdownAttachmentsInUse, cell.metadata.attachments)) {
191191
const updateMetadata: { [key: string]: any } = deepClone(cell.metadata);
192-
updateMetadata.custom.attachments = markdownAttachmentsInUse;
192+
updateMetadata.attachments = markdownAttachmentsInUse;
193193
const metadataEdit = vscode.NotebookEdit.updateCellMetadata(cell.index, updateMetadata);
194194
const workspaceEdit = new vscode.WorkspaceEdit();
195195
workspaceEdit.set(e.notebook.uri, [metadataEdit]);
@@ -229,7 +229,7 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
229229
const markdownAttachments = this.getAttachmentNames(document);
230230
if (this.checkMetadataAttachmentsExistence(activeCell.metadata)) {
231231
for (const [currFilename, attachment] of markdownAttachments) {
232-
if (!activeCell.metadata.custom.attachments[currFilename]) {
232+
if (!activeCell.metadata.attachments[currFilename]) {
233233
// no attachment reference in the metadata
234234
diagnostics.push({ name: currFilename, ranges: attachment.ranges });
235235
}
@@ -287,7 +287,7 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
287287
* @returns
288288
*/
289289
private getMetadataAttachment(metadata: { [key: string]: any }, currFilename: string): { [key: string]: any } {
290-
return metadata.custom.attachments[currFilename];
290+
return metadata.attachments[currFilename];
291291
}
292292

293293
/**
@@ -296,7 +296,7 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
296296
* @returns boolean representing the presence of any attachments
297297
*/
298298
private checkMetadataAttachmentsExistence(metadata: { [key: string]: any }): boolean {
299-
return !!(metadata.custom?.attachments);
299+
return !!(metadata.attachments);
300300
}
301301

302302
/**
@@ -311,8 +311,8 @@ export class AttachmentCleaner implements vscode.CodeActionProvider {
311311
const cellCache = documentCache.get(cellFragment) ?? new Map<string, IAttachmentData>();
312312
documentCache.set(cellFragment, cellCache);
313313

314-
for (const currFilename of Object.keys(metadata.custom.attachments)) {
315-
cellCache.set(currFilename, metadata.custom.attachments[currFilename]);
314+
for (const currFilename of Object.keys(metadata.attachments)) {
315+
cellCache.set(currFilename, metadata.attachments[currFilename]);
316316
}
317317
}
318318

‎extensions/ipynb/src/notebookImagePaste.ts

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ class CopyPasteEditProvider implements vscode.DocumentPasteEditProvider {
5151

5252
// create updated metadata for cell (prep for WorkspaceEdit)
5353
const b64string = encodeBase64(fileDataAsUint8);
54-
const startingAttachments = currentCell.metadata.custom?.attachments;
54+
const startingAttachments = currentCell.metadata.attachments;
5555
const newAttachment = buildAttachment(b64string, currentCell, filename, filetype, startingAttachments);
5656

5757
// build edits
@@ -124,13 +124,11 @@ function encodeBase64(buffer: Uint8Array, padded = true, urlSafe = false) {
124124
}
125125

126126
function buildAttachment(b64: string, cell: vscode.NotebookCell, filename: string, filetype: string, startingAttachments: any): { metadata: { [key: string]: any }; filename: string } {
127-
const outputMetadata = { ...cell.metadata };
127+
const cellMetadata = { ...cell.metadata };
128128
let tempFilename = filename + filetype;
129129

130-
if (!outputMetadata.custom) {
131-
outputMetadata['custom'] = { 'attachments': { [tempFilename]: { 'image/png': b64 } } };
132-
} else if (!outputMetadata.custom.attachments) {
133-
outputMetadata.custom['attachments'] = { [tempFilename]: { 'image/png': b64 } };
130+
if (!cellMetadata.attachments) {
131+
cellMetadata['attachments'] = { [tempFilename]: { 'image/png': b64 } };
134132
} else {
135133
for (let appendValue = 2; tempFilename in startingAttachments; appendValue++) {
136134
const objEntries = Object.entries(startingAttachments[tempFilename]);
@@ -143,10 +141,11 @@ function buildAttachment(b64: string, cell: vscode.NotebookCell, filename: strin
143141
}
144142
}
145143
}
146-
outputMetadata.custom.attachments[tempFilename] = { 'image/png': b64 };
144+
cellMetadata.attachments[tempFilename] = { 'image/png': b64 };
147145
}
146+
148147
return {
149-
metadata: outputMetadata,
148+
metadata: cellMetadata,
150149
filename: tempFilename
151150
};
152151
}

‎extensions/ipynb/src/serializers.ts

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
import type * as nbformat from '@jupyterlab/nbformat';
77
import { NotebookCell, NotebookCellData, NotebookCellKind, NotebookCellOutput } from 'vscode';
8-
import { CellMetadata, CellOutputMetadata } from './common';
8+
import { CellOutputMetadata } from './common';
99
import { textMimeTypes } from './deserializers';
1010
import { compressOutputItemStreams } from './streamCompressor';
1111

@@ -56,8 +56,14 @@ export function sortObjectPropertiesRecursively(obj: any): any {
5656
}
5757

5858
export function getCellMetadata(cell: NotebookCell | NotebookCellData) {
59-
return cell.metadata?.custom as CellMetadata | undefined;
59+
return {
60+
// it contains the cell id, and the cell metadata, along with other nb cell metadata
61+
...(cell.metadata?.custom ?? {}),
62+
// promote the cell attachments to the top level
63+
attachments: cell.metadata?.custom?.attachments ?? cell.metadata?.attachments
64+
};
6065
}
66+
6167
function createCodeCellFromNotebookCell(cell: NotebookCellData, preferredLanguage: string | undefined): nbformat.ICodeCell {
6268
const cellMetadata = getCellMetadata(cell);
6369
let metadata = cellMetadata?.metadata || {}; // This cannot be empty.
@@ -332,7 +338,7 @@ function convertOutputMimeToJupyterOutput(mime: string, value: Uint8Array) {
332338
}
333339
}
334340

335-
function createMarkdownCellFromNotebookCell(cell: NotebookCellData): nbformat.IMarkdownCell {
341+
export function createMarkdownCellFromNotebookCell(cell: NotebookCellData): nbformat.IMarkdownCell {
336342
const cellMetadata = getCellMetadata(cell);
337343
const markdownCell: any = {
338344
cell_type: 'markdown',

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

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import type * as nbformat from '@jupyterlab/nbformat';
77
import * as assert from 'assert';
88
import * as vscode from 'vscode';
99
import { jupyterCellOutputToCellOutput, jupyterNotebookModelToNotebookData } from '../deserializers';
10+
import { createMarkdownCellFromNotebookCell, getCellMetadata } from '../serializers';
1011

1112
function deepStripProperties(obj: any, props: string[]) {
1213
for (const prop in obj) {
@@ -52,6 +53,71 @@ suite('ipynb serializer', () => {
5253

5354
assert.deepStrictEqual(notebook.cells, [expectedCodeCell, expectedMarkdownCell]);
5455
});
56+
57+
58+
test('Serialize', async () => {
59+
const markdownCell = new vscode.NotebookCellData(vscode.NotebookCellKind.Markup, '# header1', 'markdown');
60+
markdownCell.metadata = {
61+
attachments: {
62+
'image.png': {
63+
'image/png': 'abc'
64+
}
65+
},
66+
custom: {
67+
id: '123',
68+
metadata: {
69+
foo: 'bar'
70+
}
71+
}
72+
};
73+
74+
const cellMetadata = getCellMetadata(markdownCell);
75+
assert.deepStrictEqual(cellMetadata, {
76+
id: '123',
77+
metadata: {
78+
foo: 'bar',
79+
},
80+
attachments: {
81+
'image.png': {
82+
'image/png': 'abc'
83+
}
84+
}
85+
});
86+
87+
const markdownCell2 = new vscode.NotebookCellData(vscode.NotebookCellKind.Markup, '# header1', 'markdown');
88+
markdownCell2.metadata = {
89+
custom: {
90+
id: '123',
91+
metadata: {
92+
foo: 'bar'
93+
},
94+
attachments: {
95+
'image.png': {
96+
'image/png': 'abc'
97+
}
98+
}
99+
}
100+
};
101+
102+
const nbMarkdownCell = createMarkdownCellFromNotebookCell(markdownCell);
103+
const nbMarkdownCell2 = createMarkdownCellFromNotebookCell(markdownCell2);
104+
assert.deepStrictEqual(nbMarkdownCell, nbMarkdownCell2);
105+
106+
assert.deepStrictEqual(nbMarkdownCell, {
107+
cell_type: 'markdown',
108+
source: ['# header1'],
109+
metadata: {
110+
foo: 'bar',
111+
},
112+
attachments: {
113+
'image.png': {
114+
'image/png': 'abc'
115+
}
116+
},
117+
id: '123'
118+
});
119+
});
120+
55121
suite('Outputs', () => {
56122
function validateCellOutputTranslation(
57123
outputs: nbformat.IOutput[],

‎src/vs/workbench/api/common/extHostTypeConverters.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1709,7 +1709,8 @@ export namespace NotebookDocumentContentOptions {
17091709
return {
17101710
transientOutputs: options?.transientOutputs ?? false,
17111711
transientCellMetadata: options?.transientCellMetadata ?? {},
1712-
transientDocumentMetadata: options?.transientDocumentMetadata ?? {}
1712+
transientDocumentMetadata: options?.transientDocumentMetadata ?? {},
1713+
cellContentMetadata: options?.cellContentMetadata ?? {}
17131714
};
17141715
}
17151716
}

‎src/vs/workbench/contrib/interactive/browser/interactive.contribution.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,8 @@ export class InteractiveDocumentContribution extends Disposable implements IWork
8787
const contentOptions = {
8888
transientOutputs: true,
8989
transientCellMetadata: {},
90-
transientDocumentMetadata: {}
90+
transientDocumentMetadata: {},
91+
cellContentMetadata: {}
9192
};
9293

9394
const controller: INotebookContentProvider = {

0 commit comments

Comments
 (0)