Port fix for transient data in notebooks. (#11168)

rchiodo · web-flow · commit 15c68db9e9ae · 2020-04-15T08:45:40.000-07:00
* Remove unwanted elements from cells when saving. (#11164) * Remove unwanted elements from cells when saving. * Fix a sonar warning * FIx another sonar warning * Update changelog
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -96,6 +96,8 @@
    ([#11035](https://github.com/Microsoft/vscode-python/issues/11035))
 1. Fix perf problems after running the interactive window for an extended period of time.
    ([#10971](https://github.com/Microsoft/vscode-python/issues/10971))
+1. Fix problem with opening a notebook in jupyter after saving in VS code.
+   ([#11151](https://github.com/Microsoft/vscode-python/issues/11151))   
 
 
 ### Code Health
diff --git a/news/2 Fixes/10971.md b/news/2 Fixes/10971.md
diff --git a/src/client/datascience/common.ts b/src/client/datascience/common.ts
@@ -1,10 +1,44 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 'use strict';
+import { nbformat } from '@jupyterlab/coreutils';
 import { Memento } from 'vscode';
+import { splitMultilineString } from '../../datascience-ui/common';
 import { noop } from '../common/utils/misc';
 import { Settings } from './constants';
 
+// Can't figure out a better way to do this. Enumerate
+// the allowed keys of different output formats.
+const dummyStreamObj: nbformat.IStream = {
+    output_type: 'stream',
+    name: 'stdout',
+    text: ''
+};
+const dummyErrorObj: nbformat.IError = {
+    output_type: 'error',
+    ename: '',
+    evalue: '',
+    traceback: ['']
+};
+const dummyDisplayObj: nbformat.IDisplayData = {
+    output_type: 'display_data',
+    data: {},
+    metadata: {}
+};
+const dummyExecuteResultObj: nbformat.IExecuteResult = {
+    output_type: 'execute_result',
+    name: '',
+    execution_count: 0,
+    data: {},
+    metadata: {}
+};
+const AllowedKeys = {
+    ['stream']: new Set(Object.keys(dummyStreamObj)),
+    ['error']: new Set(Object.keys(dummyErrorObj)),
+    ['display_data']: new Set(Object.keys(dummyDisplayObj)),
+    ['execute_result']: new Set(Object.keys(dummyExecuteResultObj))
+};
+
 export function getSavedUriList(globalState: Memento): { uri: string; time: number }[] {
     const uriList = globalState.get<{ uri: string; time: number }[]>(Settings.JupyterServerUriList);
     return uriList
@@ -23,3 +57,44 @@ export function addToUriList(globalState: Memento, uri: string, time: number) {
 
     globalState.update(Settings.JupyterServerUriList, editList).then(noop, noop);
 }
+
+function fixupOutput(output: nbformat.IOutput): nbformat.IOutput {
+    let allowedKeys: Set<string>;
+    switch (output.output_type) {
+        case 'stream':
+        case 'error':
+        case 'execute_result':
+        case 'display_data':
+            allowedKeys = AllowedKeys[output.output_type];
+            break;
+        default:
+            return output;
+    }
+    const result = { ...output };
+    for (const k of Object.keys(output)) {
+        if (!allowedKeys.has(k)) {
+            delete result[k];
+        }
+    }
+    return result;
+}
+
+export function pruneCell(cell: nbformat.ICell): nbformat.ICell {
+    // Source is usually a single string on input. Convert back to an array
+    const result = ({
+        ...cell,
+        source: splitMultilineString(cell.source)
+        // tslint:disable-next-line: no-any
+    } as any) as nbformat.ICell; // nyc (code coverage) barfs on this so just trick it.
+
+    // Remove outputs and execution_count from non code cells
+    if (result.cell_type !== 'code') {
+        delete result.outputs;
+        delete result.execution_count;
+    } else {
+        // Clean outputs from code cells
+        result.outputs = (result.outputs as nbformat.IOutput[]).map(fixupOutput);
+    }
+
+    return result;
+}
diff --git a/src/client/datascience/interactive-ipynb/nativeEditorStorage.ts b/src/client/datascience/interactive-ipynb/nativeEditorStorage.ts
@@ -20,12 +20,12 @@ import { CellState, ICell, IJupyterExecution, IJupyterKernelSpec, INotebookModel
 // tslint:disable-next-line:no-require-imports no-var-requires
 import detectIndent = require('detect-indent');
 import { sendTelemetryEvent } from '../../telemetry';
+import { pruneCell } from '../common';
 // tslint:disable-next-line:no-require-imports no-var-requires
 const debounce = require('lodash/debounce') as typeof import('lodash/debounce');
 
 const KeyPrefix = 'notebook-storage-';
 const NotebookTransferKey = 'notebook-transfered';
-
 interface INativeEditorStorageState {
     file: Uri;
     cells: ICell[];
@@ -570,23 +570,14 @@ export class NativeEditorStorage implements INotebookModel, INotebookStorage {
 
         // Reuse our original json except for the cells.
         const json = {
-            cells: cells.map((c) => this.fixupCell(c.data)),
+            cells: cells.map((c) => pruneCell(c.data)),
             metadata: this._state.notebookJson.metadata,
             nbformat: this._state.notebookJson.nbformat,
             nbformat_minor: this._state.notebookJson.nbformat_minor
         };
         return JSON.stringify(json, null, this.indentAmount);
     }
 
-    private fixupCell(cell: nbformat.ICell): nbformat.ICell {
-        // Source is usually a single string on input. Convert back to an array
-        return ({
-            ...cell,
-            source: splitMultilineString(cell.source)
-            // tslint:disable-next-line: no-any
-        } as any) as nbformat.ICell; // nyc (code coverage) barfs on this so just trick it.
-    }
-
     private getStorageKey(): string {
         return `${KeyPrefix}${this.file.toString()}`;
     }
diff --git a/src/test/datascience/datascience.unit.test.ts b/src/test/datascience/datascience.unit.test.ts
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 'use strict';
 
+import { nbformat } from '@jupyterlab/coreutils';
 import { assert } from 'chai';
 import * as sinon from 'sinon';
 import { anything, instance, mock, verify, when } from 'ts-mockito';
@@ -13,12 +14,13 @@ import { PythonSettings } from '../../client/common/configSettings';
 import { ConfigurationService } from '../../client/common/configuration/service';
 import { IConfigurationService, IPythonSettings } from '../../client/common/types';
 import { CommandRegistry } from '../../client/datascience/commands/commandRegistry';
+import { pruneCell } from '../../client/datascience/common';
 import { DataScience } from '../../client/datascience/datascience';
 import { DataScienceCodeLensProvider } from '../../client/datascience/editor-integration/codelensprovider';
 import { IDataScienceCodeLensProvider } from '../../client/datascience/types';
 
 // tslint:disable: max-func-body-length
-suite('Data Science Tests', () => {
+suite('DataScience Tests', () => {
     let dataScience: DataScience;
     let cmdManager: CommandManager;
     let codeLensProvider: IDataScienceCodeLensProvider;
@@ -75,4 +77,169 @@ suite('Data Science Tests', () => {
             assert.ok(onDidChangeActiveTextEditor.calledOnce);
         });
     });
+
+    suite('Cell pruning', () => {
+        test('Remove output and execution count from non code', () => {
+            const cell: nbformat.ICell = {
+                cell_type: 'markdown',
+                outputs: [],
+                execution_count: '23',
+                source: 'My markdown',
+                metadata: {}
+            };
+            const result = pruneCell(cell);
+            assert.equal(Object.keys(result).indexOf('outputs'), -1, 'Outputs inside markdown');
+            assert.equal(Object.keys(result).indexOf('execution_count'), -1, 'Execution count inside markdown');
+        });
+        test('Outputs dont contain extra data', () => {
+            const cell: nbformat.ICell = {
+                cell_type: 'code',
+                outputs: [
+                    {
+                        output_type: 'display_data',
+                        extra: {}
+                    }
+                ],
+                execution_count: '23',
+                source: 'My source',
+                metadata: {}
+            };
+            const result = pruneCell(cell);
+            // tslint:disable-next-line: no-any
+            assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
+            assert.equal(result.execution_count, '23', 'Output execution count removed');
+            const output = (result.outputs as nbformat.IOutput[])[0];
+            assert.equal(Object.keys(output).indexOf('extra'), -1, 'Output still has extra data');
+            assert.notEqual(Object.keys(output).indexOf('output_type'), -1, 'Output is missing output_type');
+        });
+        test('Display outputs still have their data', () => {
+            const cell: nbformat.ICell = {
+                cell_type: 'code',
+                execution_count: 2,
+                metadata: {},
+                outputs: [
+                    {
+                        output_type: 'display_data',
+                        data: {
+                            'text/plain': "Box(children=(Label(value='My label'),))",
+                            'application/vnd.jupyter.widget-view+json': {
+                                version_major: 2,
+                                version_minor: 0,
+                                model_id: '90c99248d7bb490ca132427de6d1e235'
+                            }
+                        },
+                        metadata: { bob: 'youruncle' }
+                    }
+                ],
+                source: ["line = widgets.Label('My label')\n", 'box = widgets.Box([line])\n', 'box']
+            };
+
+            const result = pruneCell(cell);
+            // tslint:disable-next-line: no-any
+            assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
+            assert.equal(result.execution_count, 2, 'Output execution count removed');
+            assert.deepEqual(result.outputs, cell.outputs, 'Outputs were modified');
+        });
+        test('Stream outputs still have their data', () => {
+            const cell: nbformat.ICell = {
+                cell_type: 'code',
+                execution_count: 2,
+                metadata: {},
+                outputs: [
+                    {
+                        output_type: 'stream',
+                        name: 'stdout',
+                        text: 'foobar'
+                    }
+                ],
+                source: ["line = widgets.Label('My label')\n", 'box = widgets.Box([line])\n', 'box']
+            };
+
+            const result = pruneCell(cell);
+            // tslint:disable-next-line: no-any
+            assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
+            assert.equal(result.execution_count, 2, 'Output execution count removed');
+            assert.deepEqual(result.outputs, cell.outputs, 'Outputs were modified');
+        });
+        test('Errors outputs still have their data', () => {
+            const cell: nbformat.ICell = {
+                cell_type: 'code',
+                execution_count: 2,
+                metadata: {},
+                outputs: [
+                    {
+                        output_type: 'error',
+                        ename: 'stdout',
+                        evalue: 'stdout is a value',
+                        traceback: ['more']
+                    }
+                ],
+                source: ["line = widgets.Label('My label')\n", 'box = widgets.Box([line])\n', 'box']
+            };
+
+            const result = pruneCell(cell);
+            // tslint:disable-next-line: no-any
+            assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
+            assert.equal(result.execution_count, 2, 'Output execution count removed');
+            assert.deepEqual(result.outputs, cell.outputs, 'Outputs were modified');
+        });
+        test('Execute result outputs still have their data', () => {
+            const cell: nbformat.ICell = {
+                cell_type: 'code',
+                execution_count: 2,
+                metadata: {},
+                outputs: [
+                    {
+                        output_type: 'execute_result',
+                        execution_count: '4',
+                        data: {
+                            'text/plain': "Box(children=(Label(value='My label'),))",
+                            'application/vnd.jupyter.widget-view+json': {
+                                version_major: 2,
+                                version_minor: 0,
+                                model_id: '90c99248d7bb490ca132427de6d1e235'
+                            }
+                        },
+                        metadata: { foo: 'bar' }
+                    }
+                ],
+                source: ["line = widgets.Label('My label')\n", 'box = widgets.Box([line])\n', 'box']
+            };
+
+            const result = pruneCell(cell);
+            // tslint:disable-next-line: no-any
+            assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
+            assert.equal(result.execution_count, 2, 'Output execution count removed');
+            assert.deepEqual(result.outputs, cell.outputs, 'Outputs were modified');
+        });
+        test('Unrecognized outputs still have their data', () => {
+            const cell: nbformat.ICell = {
+                cell_type: 'code',
+                execution_count: 2,
+                metadata: {},
+                outputs: [
+                    {
+                        output_type: 'unrecognized',
+                        execution_count: '4',
+                        data: {
+                            'text/plain': "Box(children=(Label(value='My label'),))",
+                            'application/vnd.jupyter.widget-view+json': {
+                                version_major: 2,
+                                version_minor: 0,
+                                model_id: '90c99248d7bb490ca132427de6d1e235'
+                            }
+                        },
+                        metadata: {}
+                    }
+                ],
+                source: ["line = widgets.Label('My label')\n", 'box = widgets.Box([line])\n', 'box']
+            };
+
+            const result = pruneCell(cell);
+            // tslint:disable-next-line: no-any
+            assert.equal((result.outputs as any).length, 1, 'Outputs were removed');
+            assert.equal(result.execution_count, 2, 'Output execution count removed');
+            assert.deepEqual(result.outputs, cell.outputs, 'Outputs were modified');
+        });
+    });
 });
