Skip to content

Commit bd8d32b

Browse files
rchiodorchiodo
authored
Remove unwanted elements from cells when saving. (#11164)
* Remove unwanted elements from cells when saving. * Fix a sonar warning * FIx another sonar warning
1 parent 46723ff commit bd8d32b

4 files changed

Lines changed: 246 additions & 12 deletions

File tree

‎news/2 Fixes/11151.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Fix problem with opening a notebook in jupyter after saving in VS code.

‎src/client/datascience/common.ts‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,44 @@
11
// Copyright (c) Microsoft Corporation. All rights reserved.
22
// Licensed under the MIT License.
33
'use strict';
4+
import { nbformat } from '@jupyterlab/coreutils';
45
import { Memento } from 'vscode';
6+
import { splitMultilineString } from '../../datascience-ui/common';
57
import { noop } from '../common/utils/misc';
68
import { Settings } from './constants';
79

10+
// Can't figure out a better way to do this. Enumerate
11+
// the allowed keys of different output formats.
12+
const dummyStreamObj: nbformat.IStream = {
13+
output_type: 'stream',
14+
name: 'stdout',
15+
text: ''
16+
};
17+
const dummyErrorObj: nbformat.IError = {
18+
output_type: 'error',
19+
ename: '',
20+
evalue: '',
21+
traceback: ['']
22+
};
23+
const dummyDisplayObj: nbformat.IDisplayData = {
24+
output_type: 'display_data',
25+
data: {},
26+
metadata: {}
27+
};
28+
const dummyExecuteResultObj: nbformat.IExecuteResult = {
29+
output_type: 'execute_result',
30+
name: '',
31+
execution_count: 0,
32+
data: {},
33+
metadata: {}
34+
};
35+
const AllowedKeys = {
36+
['stream']: new Set(Object.keys(dummyStreamObj)),
37+
['error']: new Set(Object.keys(dummyErrorObj)),
38+
['display_data']: new Set(Object.keys(dummyDisplayObj)),
39+
['execute_result']: new Set(Object.keys(dummyExecuteResultObj))
40+
};
41+
842
export function getSavedUriList(globalState: Memento): { uri: string; time: number }[] {
943
const uriList = globalState.get<{ uri: string; time: number }[]>(Settings.JupyterServerUriList);
1044
return uriList
@@ -23,3 +57,44 @@ export function addToUriList(globalState: Memento, uri: string, time: number) {
2357

2458
globalState.update(Settings.JupyterServerUriList, editList).then(noop, noop);
2559
}
60+
61+
function fixupOutput(output: nbformat.IOutput): nbformat.IOutput {
62+
let allowedKeys: Set<string>;
63+
switch (output.output_type) {
64+
case 'stream':
65+
case 'error':
66+
case 'execute_result':
67+
case 'display_data':
68+
allowedKeys = AllowedKeys[output.output_type];
69+
break;
70+
default:
71+
return output;
72+
}
73+
const result = { ...output };
74+
for (const k of Object.keys(output)) {
75+
if (!allowedKeys.has(k)) {
76+
delete result[k];
77+
}
78+
}
79+
return result;
80+
}
81+
82+
export function pruneCell(cell: nbformat.ICell): nbformat.ICell {
83+
// Source is usually a single string on input. Convert back to an array
84+
const result = ({
85+
...cell,
86+
source: splitMultilineString(cell.source)
87+
// tslint:disable-next-line: no-any
88+
} as any) as nbformat.ICell; // nyc (code coverage) barfs on this so just trick it.
89+
90+
// Remove outputs and execution_count from non code cells
91+
if (result.cell_type !== 'code') {
92+
delete result.outputs;
93+
delete result.execution_count;
94+
} else {
95+
// Clean outputs from code cells
96+
result.outputs = (result.outputs as nbformat.IOutput[]).map(fixupOutput);
97+
}
98+
99+
return result;
100+
}

‎src/client/datascience/interactive-ipynb/nativeEditorStorage.ts‎

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,12 @@ import { CellState, ICell, IJupyterExecution, IJupyterKernelSpec, INotebookModel
2020
// tslint:disable-next-line:no-require-imports no-var-requires
2121
import detectIndent = require('detect-indent');
2222
import { sendTelemetryEvent } from '../../telemetry';
23+
import { pruneCell } from '../common';
2324
// tslint:disable-next-line:no-require-imports no-var-requires
2425
const debounce = require('lodash/debounce') as typeof import('lodash/debounce');
2526

2627
const KeyPrefix = 'notebook-storage-';
2728
const NotebookTransferKey = 'notebook-transfered';
28-
2929
interface INativeEditorStorageState {
3030
file: Uri;
3131
cells: ICell[];
@@ -570,23 +570,14 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
570570

571571
// Reuse our original json except for the cells.
572572
const json = {
573-
cells: cells.map((c) => this.fixupCell(c.data)),
573+
cells: cells.map((c) => pruneCell(c.data)),
574574
metadata: this._state.notebookJson.metadata,
575575
nbformat: this._state.notebookJson.nbformat,
576576
nbformat_minor: this._state.notebookJson.nbformat_minor
577577
};
578578
return JSON.stringify(json, null, this.indentAmount);
579579
}
580580

581-
private fixupCell(cell: nbformat.ICell): nbformat.ICell {
582-
// Source is usually a single string on input. Convert back to an array
583-
return ({
584-
...cell,
585-
source: splitMultilineString(cell.source)
586-
// tslint:disable-next-line: no-any
587-
} as any) as nbformat.ICell; // nyc (code coverage) barfs on this so just trick it.
588-
}
589-
590581
private getStorageKey(): string {
591582
return `${KeyPrefix}${this.file.toString()}`;
592583
}

‎src/test/datascience/datascience.unit.test.ts‎

Lines changed: 168 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// Licensed under the MIT License.
33
'use strict';
44

5+
import { nbformat } from '@jupyterlab/coreutils';
56
import { assert } from 'chai';
67
import * as sinon from 'sinon';
78
import { anything, instance, mock, verify, when } from 'ts-mockito';
@@ -13,12 +14,13 @@ import { PythonSettings } from '../../client/common/configSettings';
1314
import { ConfigurationService } from '../../client/common/configuration/service';
1415
import { IConfigurationService, IPythonSettings } from '../../client/common/types';
1516
import { CommandRegistry } from '../../client/datascience/commands/commandRegistry';
17+
import { pruneCell } from '../../client/datascience/common';
1618
import { DataScience } from '../../client/datascience/datascience';
1719
import { DataScienceCodeLensProvider } from '../../client/datascience/editor-integration/codelensprovider';
1820
import { IDataScienceCodeLensProvider } from '../../client/datascience/types';
1921

2022
// tslint:disable: max-func-body-length
21-
suite('Data Science Tests', () => {
23+
suite('DataScience Tests', () => {
2224
let dataScience: DataScience;
2325
let cmdManager: CommandManager;
2426
let codeLensProvider: IDataScienceCodeLensProvider;
@@ -75,4 +77,169 @@ suite('Data Science Tests', () => {
7577
assert.ok(onDidChangeActiveTextEditor.calledOnce);
7678
});
7779
});
80+
81+
suite('Cell pruning', () => {
82+
test('Remove output and execution count from non code', () => {
83+
const cell: nbformat.ICell = {
84+
cell_type: 'markdown',
85+
outputs: [],
86+
execution_count: '23',
87+
source: 'My markdown',
88+
metadata: {}
89+
};
90+
const result = pruneCell(cell);
91+
assert.equal(Object.keys(result).indexOf('outputs'), -1, 'Outputs inside markdown');
92+
assert.equal(Object.keys(result).indexOf('execution_count'), -1, 'Execution count inside markdown');
93+
});
94+
test('Outputs dont contain extra data', () => {
95+
const cell: nbformat.ICell = {
96+
cell_type: 'code',
97+
outputs: [
98+
{
99+
output_type: 'display_data',
100+
extra: {}
101+
}
102+
],
103+
execution_count: '23',
104+
source: 'My source',
105+
metadata: {}
106+
};
107+
const result = pruneCell(cell);
108+
// tslint:disable-next-line: no-any
109+
assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
110+
assert.equal(result.execution_count, '23', 'Output execution count removed');
111+
const output = (result.outputs as nbformat.IOutput[])[0];
112+
assert.equal(Object.keys(output).indexOf('extra'), -1, 'Output still has extra data');
113+
assert.notEqual(Object.keys(output).indexOf('output_type'), -1, 'Output is missing output_type');
114+
});
115+
test('Display outputs still have their data', () => {
116+
const cell: nbformat.ICell = {
117+
cell_type: 'code',
118+
execution_count: 2,
119+
metadata: {},
120+
outputs: [
121+
{
122+
output_type: 'display_data',
123+
data: {
124+
'text/plain': "Box(children=(Label(value='My label'),))",
125+
'application/vnd.jupyter.widget-view+json': {
126+
version_major: 2,
127+
version_minor: 0,
128+
model_id: '90c99248d7bb490ca132427de6d1e235'
129+
}
130+
},
131+
metadata: { bob: 'youruncle' }
132+
}
133+
],
134+
source: ["line = widgets.Label('My label')\n", 'box = widgets.Box([line])\n', 'box']
135+
};
136+
137+
const result = pruneCell(cell);
138+
// tslint:disable-next-line: no-any
139+
assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
140+
assert.equal(result.execution_count, 2, 'Output execution count removed');
141+
assert.deepEqual(result.outputs, cell.outputs, 'Outputs were modified');
142+
});
143+
test('Stream outputs still have their data', () => {
144+
const cell: nbformat.ICell = {
145+
cell_type: 'code',
146+
execution_count: 2,
147+
metadata: {},
148+
outputs: [
149+
{
150+
output_type: 'stream',
151+
name: 'stdout',
152+
text: 'foobar'
153+
}
154+
],
155+
source: ["line = widgets.Label('My label')\n", 'box = widgets.Box([line])\n', 'box']
156+
};
157+
158+
const result = pruneCell(cell);
159+
// tslint:disable-next-line: no-any
160+
assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
161+
assert.equal(result.execution_count, 2, 'Output execution count removed');
162+
assert.deepEqual(result.outputs, cell.outputs, 'Outputs were modified');
163+
});
164+
test('Errors outputs still have their data', () => {
165+
const cell: nbformat.ICell = {
166+
cell_type: 'code',
167+
execution_count: 2,
168+
metadata: {},
169+
outputs: [
170+
{
171+
output_type: 'error',
172+
ename: 'stdout',
173+
evalue: 'stdout is a value',
174+
traceback: ['more']
175+
}
176+
],
177+
source: ["line = widgets.Label('My label')\n", 'box = widgets.Box([line])\n', 'box']
178+
};
179+
180+
const result = pruneCell(cell);
181+
// tslint:disable-next-line: no-any
182+
assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
183+
assert.equal(result.execution_count, 2, 'Output execution count removed');
184+
assert.deepEqual(result.outputs, cell.outputs, 'Outputs were modified');
185+
});
186+
test('Execute result outputs still have their data', () => {
187+
const cell: nbformat.ICell = {
188+
cell_type: 'code',
189+
execution_count: 2,
190+
metadata: {},
191+
outputs: [
192+
{
193+
output_type: 'execute_result',
194+
execution_count: '4',
195+
data: {
196+
'text/plain': "Box(children=(Label(value='My label'),))",
197+
'application/vnd.jupyter.widget-view+json': {
198+
version_major: 2,
199+
version_minor: 0,
200+
model_id: '90c99248d7bb490ca132427de6d1e235'
201+
}
202+
},
203+
metadata: { foo: 'bar' }
204+
}
205+
],
206+
source: ["line = widgets.Label('My label')\n", 'box = widgets.Box([line])\n', 'box']
207+
};
208+
209+
const result = pruneCell(cell);
210+
// tslint:disable-next-line: no-any
211+
assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
212+
assert.equal(result.execution_count, 2, 'Output execution count removed');
213+
assert.deepEqual(result.outputs, cell.outputs, 'Outputs were modified');
214+
});
215+
test('Unrecognized outputs still have their data', () => {
216+
const cell: nbformat.ICell = {
217+
cell_type: 'code',
218+
execution_count: 2,
219+
metadata: {},
220+
outputs: [
221+
{
222+
output_type: 'unrecognized',
223+
execution_count: '4',
224+
data: {
225+
'text/plain': "Box(children=(Label(value='My label'),))",
226+
'application/vnd.jupyter.widget-view+json': {
227+
version_major: 2,
228+
version_minor: 0,
229+
model_id: '90c99248d7bb490ca132427de6d1e235'
230+
}
231+
},
232+
metadata: {}
233+
}
234+
],
235+
source: ["line = widgets.Label('My label')\n", 'box = widgets.Box([line])\n', 'box']
236+
};
237+
238+
const result = pruneCell(cell);
239+
// tslint:disable-next-line: no-any
240+
assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
241+
assert.equal(result.execution_count, 2, 'Output execution count removed');
242+
assert.deepEqual(result.outputs, cell.outputs, 'Outputs were modified');
243+
});
244+
});
78245
});

0 commit comments

Comments
 (0)