Port prune fix from main to release (#14243)

* Remove unneeded cell keys when exporting (#14241)

* Remove transient output when exporting from the interactive window

* Add news entry

* Update changelog

Author: rchiodo

news/2 Fixes/14213.md

@@ -38,7 +38,7 @@ const dummyExecuteResultObj: nbformat.IExecuteResult = {
     data: {},
     metadata: {}
 };
-const AllowedKeys = {
+export const AllowedCellOutputKeys = {
     ['stream']: new Set(Object.keys(dummyStreamObj)),
     ['error']: new Set(Object.keys(dummyErrorObj)),
     ['display_data']: new Set(Object.keys(dummyDisplayObj)),
@@ -73,7 +73,7 @@ function fixupOutput(output: nbformat.IOutput): nbformat.IOutput {
         case 'error':
         case 'execute_result':
         case 'display_data':
-            allowedKeys = AllowedKeys[output.output_type];
+            allowedKeys = AllowedCellOutputKeys[output.output_type];
             break;
         default:
             return output;

src/client/datascience/common.ts

@@ -17,6 +17,7 @@ import { IConfigurationService } from '../../common/types';
 import * as localize from '../../common/utils/localize';
 import { noop } from '../../common/utils/misc';
 import { CellMatcher } from '../cellMatcher';
+import { pruneCell } from '../common';
 import { CodeSnippets, Identifiers } from '../constants';
 import {
     CellState,
@@ -235,9 +236,11 @@ export class JupyterExporter implements INotebookExporter {
     };
 
     private pruneCell = (cell: ICell, cellMatcher: CellMatcher): nbformat.IBaseCell => {
+        // Prune with the common pruning function first.
+        const copy = pruneCell({ ...cell.data });
+
         // Remove the #%% of the top of the source if there is any. We don't need
         // this to end up in the exported ipynb file.
-        const copy = { ...cell.data };
         copy.source = this.pruneSource(cell.data.source, cellMatcher);
         return copy;
     };

src/client/datascience/jupyter/jupyterExporter.ts

@@ -9,6 +9,7 @@ import * as path from 'path';
 import * as TypeMoq from 'typemoq';
 import { Disposable, Memento, Selection, TextDocument, TextEditor, Uri } from 'vscode';
 
+import { nbformat } from '@jupyterlab/coreutils';
 import { ReactWrapper } from 'enzyme';
 import { anything, when } from 'ts-mockito';
 import { IApplicationShell, IDocumentManager } from '../../client/common/application/types';
@@ -17,11 +18,12 @@ import { createDeferred, sleep, waitForPromise } from '../../client/common/utils
 import { noop } from '../../client/common/utils/misc';
 import { EXTENSION_ROOT_DIR } from '../../client/constants';
 import { generateCellsFromDocument } from '../../client/datascience/cellFactory';
+import { AllowedCellOutputKeys } from '../../client/datascience/common';
 import { EditorContexts } from '../../client/datascience/constants';
 import { InteractiveWindowMessages } from '../../client/datascience/interactive-common/interactiveWindowTypes';
 import { InteractiveWindow } from '../../client/datascience/interactive-window/interactiveWindow';
 import { AskedForPerFileSettingKey } from '../../client/datascience/interactive-window/interactiveWindowProvider';
-import { IInteractiveWindowProvider } from '../../client/datascience/types';
+import { IDataScienceFileSystem, IInteractiveWindowProvider } from '../../client/datascience/types';
 import { IInterpreterService } from '../../client/interpreter/contracts';
 import { concatMultilineString } from '../../datascience-ui/common';
 import { InteractivePanel } from '../../datascience-ui/history-react/interactivePanel';
@@ -635,53 +637,81 @@ for i in range(0, 100):
                     return;
                 }
             };
-            let exportCalled = false;
-            const appShell = TypeMoq.Mock.ofType<IApplicationShell>();
-            appShell
-                .setup((a) => a.showErrorMessage(TypeMoq.It.isAnyString()))
-                .returns((e) => {
-                    throw e;
-                });
-            appShell
-                .setup((a) => a.showInformationMessage(TypeMoq.It.isAny(), TypeMoq.It.isAny()))
-                .returns(() => Promise.resolve(''));
-            appShell
-                .setup((a) => a.showSaveDialog(TypeMoq.It.isAny()))
-                .returns(() => {
-                    exportCalled = true;
-                    return Promise.resolve(undefined);
-                });
-            appShell.setup((a) => a.setStatusBarMessage(TypeMoq.It.isAny())).returns(() => dummyDisposable);
-            ioc.serviceManager.rebindInstance<IApplicationShell>(IApplicationShell, appShell.object);
-
-            // Make sure to create the interactive window after the rebind or it gets the wrong application shell.
-            await addCode(ioc, 'a=1\na');
-            const { window, mount } = await getOrCreateInteractiveWindow(ioc);
-
-            // Export should cause exportCalled to change to true
-            const exportPromise = mount.waitForMessage(InteractiveWindowMessages.ReturnAllCells);
-            window.exportCells();
-            await exportPromise;
-            await sleep(100); // Give time for appshell to come up
-            assert.equal(exportCalled, true, 'Export is not being called during export');
+            const dsfs = ioc.get<IDataScienceFileSystem>(IDataScienceFileSystem);
+            const tf = await dsfs.createTemporaryLocalFile('.ipynb');
+            try {
+                let exportCalled = false;
+                const appShell = TypeMoq.Mock.ofType<IApplicationShell>();
+                appShell
+                    .setup((a) => a.showErrorMessage(TypeMoq.It.isAnyString()))
+                    .returns((e) => {
+                        throw e;
+                    });
+                appShell
+                    .setup((a) => a.showInformationMessage(TypeMoq.It.isAny(), TypeMoq.It.isAny()))
+                    .returns(() => Promise.resolve(''));
+                appShell
+                    .setup((a) => a.showSaveDialog(TypeMoq.It.isAny()))
+                    .returns(() => {
+                        exportCalled = true;
+                        return Promise.resolve(Uri.file(tf.filePath));
+                    });
+                appShell.setup((a) => a.setStatusBarMessage(TypeMoq.It.isAny())).returns(() => dummyDisposable);
+                ioc.serviceManager.rebindInstance<IApplicationShell>(IApplicationShell, appShell.object);
+                const exportCode = `
+for i in range(100):
+    time.sleep(0.1)
+    raise Exception('test')
+`;
 
-            // Remove the cell
-            const exportButton = findButton(mount.wrapper, InteractivePanel, 6);
-            const undo = findButton(mount.wrapper, InteractivePanel, 2);
+                // Make sure to create the interactive window after the rebind or it gets the wrong application shell.
+                addMockData(ioc, exportCode, 'NameError', 'type/error', 'error', [
+                    '\u001b[1;31m---------------------------------------------------------------------------\u001b[0m',
+                    '\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)',
+                    "\u001b[1;32md:\\Source\\Testing_3\\manualTestFile.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m100\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m     \u001b[0mtime\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msleep\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m0.1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      3\u001b[0m     \u001b[1;32mraise\u001b[0m \u001b[0mException\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'test'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
+                    "\u001b[1;31mNameError\u001b[0m: name 'time' is not defined"
+                ]);
+                await addCode(ioc, exportCode);
+                const { window, mount } = await getOrCreateInteractiveWindow(ioc);
 
-            // Now verify if we undo, we have no cells
-            const afterUndo = await getInteractiveCellResults(ioc, () => {
-                undo!.simulate('click');
-                return Promise.resolve();
-            });
+                // Export should cause exportCalled to change to true
+                const exportPromise = mount.waitForMessage(InteractiveWindowMessages.ReturnAllCells);
+                window.exportCells();
+                await exportPromise;
+                await sleep(100); // Give time for appshell to come up
+                assert.equal(exportCalled, true, 'Export is not being called during export');
+
+                // Read file contents into a jupyter structure. Make sure we have only the expected values
+                const contents = await dsfs.readLocalFile(tf.filePath);
+                const struct = JSON.parse(contents) as nbformat.INotebookContent;
+                assert.strictEqual(struct.cells.length, 1, 'Wrong number of cells');
+                const outputs = struct.cells[0].outputs as nbformat.IOutput[];
+                assert.strictEqual(outputs.length, 1, 'Not correct number of outputs');
+                assert.strictEqual(outputs[0].output_type, 'error', 'Error not found');
+                const allowedKeys = [...AllowedCellOutputKeys.error];
+                const actualKeys = Object.keys(outputs[0]);
+                assert.deepStrictEqual(allowedKeys, actualKeys, 'Invalid keys in output');
+
+                // Remove the cell
+                const exportButton = findButton(mount.wrapper, InteractivePanel, 6);
+                const undo = findButton(mount.wrapper, InteractivePanel, 2);
+
+                // Now verify if we undo, we have no cells
+                const afterUndo = await getInteractiveCellResults(ioc, () => {
+                    undo!.simulate('click');
+                    return Promise.resolve();
+                });
 
-            assert.equal(afterUndo.length, 1, 'Undo should remove cells');
+                assert.equal(afterUndo.length, 1, 'Undo should remove cells');
 
-            // Then verify we cannot click the button (it should be disabled)
-            exportCalled = false;
-            exportButton!.simulate('click');
-            await sleep(100);
-            assert.equal(exportCalled, false, 'Export should not be called when no cells visible');
+                // Then verify we cannot click the button (it should be disabled)
+                exportCalled = false;
+                exportButton!.simulate('click');
+                await sleep(100);
+                assert.equal(exportCalled, false, 'Export should not be called when no cells visible');
+            } finally {
+                tf.dispose();
+            }
         },
         () => {
             return ioc;

src/test/datascience/interactiveWindow.functional.test.tsx

@@ -289,13 +289,13 @@ export class MockJupyterManager implements IJupyterSessionManager {
         this.makeActive(interpreter);
     }
 
-    public addError(code: string, message: string) {
+    public addError(code: string, message: string, traceback?: string[]) {
         // Turn the message into an nbformat.IError
         const result: nbformat.IError = {
             output_type: 'error',
             ename: message,
             evalue: message,
-            traceback: [message]
+            traceback: traceback ? traceback : [message]
         };
 
         this.addCell(code, result);

src/test/datascience/mockJupyterManager.ts

@@ -89,11 +89,12 @@ export function addMockData(
     code: string,
     result: string | number | undefined | string[],
     mimeType?: string | string[],
-    cellType?: string
+    cellType?: string,
+    traceback?: string[]
 ) {
     if (ioc.mockJupyter) {
         if (cellType && cellType === 'error') {
-            ioc.mockJupyter.addError(code, result ? result.toString() : '');
+            ioc.mockJupyter.addError(code, result ? result.toString() : '', traceback);
         } else {
             if (result) {
                 ioc.mockJupyter.addCell(code, result, mimeType);