Update other cells in cell execution (#14240)

DonJayamanne · web-flow · commit 1626cf959460 · 2020-10-02T16:27:48.000-07:00
* Update other cells in cell execution

* Add tests
diff --git a/src/client/datascience/baseJupyterSession.ts b/src/client/datascience/baseJupyterSession.ts
@@ -68,9 +68,6 @@ export abstract class BaseJupyterSession implements IJupyterSession {
     public get isConnected(): boolean {
         return this.connected;
     }
-    public get onIoPubMessage() {
-        return this.ioPubEventEmitter.event;
-    }
     protected onStatusChangedEvent: EventEmitter<ServerStatus> = new EventEmitter<ServerStatus>();
     protected statusHandler: Slot<ISessionWithSocket, Kernel.Status>;
     protected connected: boolean = false;
diff --git a/src/client/datascience/jupyter/kernels/cellExecution.ts b/src/client/datascience/jupyter/kernels/cellExecution.ts
@@ -18,7 +18,11 @@ import { noop } from '../../../common/utils/misc';
 import { StopWatch } from '../../../common/utils/stopWatch';
 import { sendTelemetryEvent } from '../../../telemetry';
 import { Telemetry } from '../../constants';
-import { updateCellExecutionCount, updateCellWithErrorStatus } from '../../notebook/helpers/executionHelpers';
+import {
+    handleUpdateDisplayDataMessage,
+    updateCellExecutionCount,
+    updateCellWithErrorStatus
+} from '../../notebook/helpers/executionHelpers';
 import {
     cellOutputToVSCCellOutput,
     clearCellForExecution,
@@ -365,12 +369,12 @@ export class CellExecution {
                 await this.handleStreamMessage(msg as KernelMessage.IStreamMsg, clearState);
             } else if (jupyterLab.KernelMessage.isDisplayDataMsg(msg)) {
                 await this.handleDisplayData(msg as KernelMessage.IDisplayDataMsg, clearState);
+            } else if (jupyterLab.KernelMessage.isUpdateDisplayDataMsg(msg)) {
+                await handleUpdateDisplayDataMessage(msg, this.editor);
             } else if (jupyterLab.KernelMessage.isClearOutputMsg(msg)) {
                 await this.handleClearOutput(msg as KernelMessage.IClearOutputMsg, clearState);
             } else if (jupyterLab.KernelMessage.isErrorMsg(msg)) {
                 await this.handleError(msg as KernelMessage.IErrorMsg, clearState);
-            } else if (jupyterLab.KernelMessage.isUpdateDisplayDataMsg(msg)) {
-                // Noop.
             } else if (jupyterLab.KernelMessage.isCommOpenMsg(msg)) {
                 // Noop.
             } else if (jupyterLab.KernelMessage.isCommMsgMsg(msg)) {
diff --git a/src/client/datascience/jupyter/kernels/kernelExecution.ts b/src/client/datascience/jupyter/kernels/kernelExecution.ts
@@ -3,15 +3,13 @@
 
 'use strict';
 
-import { KernelMessage } from '@jupyterlab/services';
 import { NotebookCell, NotebookCellRunState, NotebookDocument } from 'vscode';
 import { IApplicationShell, ICommandManager, IVSCodeNotebook } from '../../../common/application/types';
 import { IDisposable } from '../../../common/types';
 import { noop } from '../../../common/utils/misc';
 import { IInterpreterService } from '../../../interpreter/contracts';
 import { captureTelemetry } from '../../../telemetry';
 import { Commands, Telemetry, VSCodeNativeTelemetry } from '../../constants';
-import { handleUpdateDisplayDataMessage } from '../../notebook/helpers/executionHelpers';
 import { MultiCancellationTokenSource } from '../../notebook/helpers/multiCancellationToken';
 import { IDataScienceErrorHandler, INotebook, INotebookEditorProvider } from '../../types';
 import { CellExecution, CellExecutionFactory } from './cellExecution';
@@ -161,15 +159,6 @@ export class KernelExecution implements IDisposable {
         return kernel;
     }
 
-    private async onIoPubMessage(document: NotebookDocument, msg: KernelMessage.IIOPubMessage) {
-        // tslint:disable-next-line:no-require-imports
-        const jupyterLab = require('@jupyterlab/services') as typeof import('@jupyterlab/services');
-        const editor = this.vscNotebook.notebookEditors.find((e) => e.document === document);
-        if (jupyterLab.KernelMessage.isUpdateDisplayDataMsg(msg) && editor) {
-            await handleUpdateDisplayDataMessage(msg, editor);
-        }
-    }
-
     private async executeIndividualCell(
         kernelPromise: Promise<IKernel>,
         cellExecution: CellExecution
@@ -178,20 +167,9 @@ export class KernelExecution implements IDisposable {
             throw new Error('No notebook object');
         }
 
-        // Register for IO pub messages
-        const ioRegistration = this.notebook.session.onIoPubMessage(
-            this.onIoPubMessage.bind(this, cellExecution.cell.notebook)
-        );
         cellExecution.token.onCancellationRequested(
-            () => {
-                ioRegistration.dispose();
-                if (cellExecution.completed) {
-                    return;
-                }
-
-                // Interrupt kernel only if we need to cancel a cell execution.
-                this.commandManager.executeCommand(Commands.NotebookEditorInterruptKernel).then(noop, noop);
-            },
+            // Interrupt kernel only if we need to cancel a cell execution.
+            () => this.commandManager.executeCommand(Commands.NotebookEditorInterruptKernel).then(noop, noop),
             this,
             this.disposables
         );
@@ -200,11 +178,7 @@ export class KernelExecution implements IDisposable {
         await cellExecution.start(kernelPromise, this.notebook);
 
         // The result promise will resolve when complete.
-        try {
-            return await cellExecution.result;
-        } finally {
-            ioRegistration.dispose();
-        }
+        return cellExecution.result;
     }
 
     private async validateKernel(document: NotebookDocument): Promise<void> {
diff --git a/src/client/datascience/types.ts b/src/client/datascience/types.ts
@@ -320,7 +320,6 @@ export interface IJupyterPasswordConnect {
 export const IJupyterSession = Symbol('IJupyterSession');
 export interface IJupyterSession extends IAsyncDisposable {
     onSessionStatusChanged: Event<ServerStatus>;
-    onIoPubMessage: Event<KernelMessage.IIOPubMessage>;
     readonly status: ServerStatus;
     readonly workingDirectory: string;
     readonly kernelSocket: Observable<KernelSocketInformation | undefined>;
diff --git a/src/test/datascience/mockJupyterSession.ts b/src/test/datascience/mockJupyterSession.ts
@@ -25,9 +25,6 @@ export class MockJupyterSession implements IJupyterSession {
     private dict: Record<string, ICell>;
     private restartedEvent: EventEmitter<void> = new EventEmitter<void>();
     private onStatusChangedEvent: EventEmitter<ServerStatus> = new EventEmitter<ServerStatus>();
-    private onIoPubMessageEvent: EventEmitter<KernelMessage.IIOPubMessage> = new EventEmitter<
-        KernelMessage.IIOPubMessage
-    >();
     private timedelay: number;
     private executionCount: number = 0;
     private outstandingRequestTokenSources: CancellationTokenSource[] = [];
@@ -58,10 +55,6 @@ export class MockJupyterSession implements IJupyterSession {
         }
         return this.onStatusChangedEvent.event;
     }
-    public get onIoPubMessage(): Event<KernelMessage.IIOPubMessage> {
-        return this.onIoPubMessageEvent.event;
-    }
-
     public get status(): ServerStatus {
         return this._status;
     }
diff --git a/src/test/datascience/notebook/executionService.ds.test.ts b/src/test/datascience/notebook/executionService.ds.test.ts
@@ -25,7 +25,7 @@ import {
     deleteAllCellsAndWait,
     executeActiveDocument,
     executeCell,
-    insertPythonCellAndWait,
+    insertPythonCell,
     startJupyter,
     trustAllNotebooks
 } from './helper';
@@ -56,7 +56,7 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
     setup(deleteAllCellsAndWait);
     suiteTeardown(() => closeNotebooksAndCleanUpAfterTests(disposables));
     test('Execute cell using VSCode Kernel', async () => {
-        await insertPythonCellAndWait('print("Hello World")');
+        await insertPythonCell('print("Hello World")');
         const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
 
         await executeCell(cell);
@@ -69,7 +69,7 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         );
     });
     test('Executed events are triggered', async () => {
-        await insertPythonCellAndWait('print("Hello World")');
+        await insertPythonCell('print("Hello World")');
         const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
 
         const executed = createEventHandler(editorProvider.activeEditor!, 'executed', disposables);
@@ -87,7 +87,7 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         await codeExecuted.assertFired(1_000);
     });
     test('Empty cell will not get executed', async () => {
-        await insertPythonCellAndWait('');
+        await insertPythonCell('');
         const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
         await executeCell(cell);
 
@@ -96,8 +96,8 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         assert.isUndefined(cell?.metadata.runState);
     });
     test('Empty cells will not get executed when running whole document', async () => {
-        await insertPythonCellAndWait('');
-        await insertPythonCellAndWait('print("Hello World")');
+        await insertPythonCell('');
+        await insertPythonCell('print("Hello World")');
         const cells = vscodeNotebook.activeNotebookEditor?.document.cells!;
 
         await executeActiveDocument();
@@ -111,7 +111,7 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         assert.isUndefined(cells[0].metadata.runState);
     });
     test('Verify Cell output, execution count and status', async () => {
-        await insertPythonCellAndWait('print("Hello World")');
+        await insertPythonCell('print("Hello World")');
         const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
 
         await executeActiveDocument();
@@ -130,8 +130,8 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         assert.ok(cell.metadata.executionOrder, 'Execution count should be > 0');
     });
     test('Verify multiple cells get executed', async () => {
-        await insertPythonCellAndWait('print("Foo Bar")');
-        await insertPythonCellAndWait('print("Hello World")');
+        await insertPythonCell('print("Foo Bar")');
+        await insertPythonCell('print("Hello World")');
         const cells = vscodeNotebook.activeNotebookEditor?.document.cells!;
 
         await executeActiveDocument();
@@ -153,7 +153,7 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         assert.equal(cells[1].metadata.executionOrder! - 1, cells[0].metadata.executionOrder!);
     });
     test('Verify metadata for successfully executed cell', async () => {
-        await insertPythonCellAndWait('print("Foo Bar")');
+        await insertPythonCell('print("Foo Bar")');
         const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
 
         await executeActiveDocument();
@@ -172,7 +172,7 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         assert.equal(cell.metadata.statusMessage, '', 'Incorrect Status message');
     });
     test('Verify output & metadata for executed cell with errors', async () => {
-        await insertPythonCellAndWait('print(abcd)');
+        await insertPythonCell('print(abcd)');
         const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
 
         await executeActiveDocument();
@@ -198,9 +198,9 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         assert.include(cell.metadata.statusMessage!, 'abcd', 'Must contain error message');
     });
     test('Updating display data', async () => {
-        await insertPythonCellAndWait('from IPython.display import Markdown\n');
-        await insertPythonCellAndWait('dh = display(display_id=True)\n');
-        await insertPythonCellAndWait('dh.update(Markdown("foo"))\n');
+        await insertPythonCell('from IPython.display import Markdown\n');
+        await insertPythonCell('dh = display(display_id=True)\n');
+        await insertPythonCell('dh.update(Markdown("foo"))\n');
         const displayCell = vscodeNotebook.activeNotebookEditor?.document.cells![1]!;
         const updateCell = vscodeNotebook.activeNotebookEditor?.document.cells![2]!;
 
@@ -225,7 +225,7 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         // Assume you are executing a cell that prints numbers 1-100.
         // When printing number 50, you click clear.
         // Cell output should now start printing output from 51 onwards, & not 1.
-        await insertPythonCellAndWait(
+        await insertPythonCell(
             dedent`
                     print("Start")
                     import time
@@ -277,7 +277,7 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         // Assume you are executing a cell that prints numbers 1-100.
         // When printing number 50, you click clear.
         // Cell output should now start printing output from 51 onwards, & not 1.
-        await insertPythonCellAndWait(
+        await insertPythonCell(
             dedent`
                 from IPython.display import display, clear_output
                 import time
@@ -315,7 +315,7 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         // Assume you are executing a cell that prints numbers 1-100.
         // When printing number 50, you click clear.
         // Cell output should now start printing output from 51 onwards, & not 1.
-        await insertPythonCellAndWait(
+        await insertPythonCell(
             dedent`
                     print("Start")
                     import time
@@ -344,14 +344,14 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         );
     });
     test('Verify escaping of output', async () => {
-        await insertPythonCellAndWait('1');
-        await insertPythonCellAndWait(dedent`
+        await insertPythonCell('1');
+        await insertPythonCell(dedent`
                                             a="<a href=f>"
                                             a`);
-        await insertPythonCellAndWait(dedent`
+        await insertPythonCell(dedent`
                                             a="<a href=f>"
                                             print(a)`);
-        await insertPythonCellAndWait('raise Exception("<whatever>")');
+        await insertPythonCell('raise Exception("<whatever>")');
         const cells = vscodeNotebook.activeNotebookEditor?.document.cells!;
 
         await executeActiveDocument();
@@ -391,4 +391,44 @@ suite('DataScience - VSCode Notebook - (Execution) (slow)', function () {
         assert.isNotEmpty(errorOutput.traceback, 'Incorrect traceback');
         assert.include(errorOutput.traceback.join(''), '<whatever>');
     });
+    test('Verify display updates', async () => {
+        await insertPythonCell('from IPython.display import Markdown', 0);
+        await insertPythonCell('dh = display(Markdown("foo"), display_id=True)', 1);
+        let cells = vscodeNotebook.activeNotebookEditor?.document.cells!;
+
+        await executeActiveDocument();
+        await waitForCondition(
+            async () => assertHasExecutionCompletedSuccessfully(cells[1]),
+            15_000,
+            'Cell did not get executed'
+        );
+
+        assert.equal(cells[0].outputs.length, 0, 'Incorrect number of output');
+        assert.equal(cells[1].outputs.length, 1, 'Incorrect number of output');
+        assert.equal(cells[1].outputs[0].outputKind, vscodeNotebookEnums.CellOutputKind.Rich, 'Incorrect output type');
+        assert.equal((cells[1].outputs[0] as CellDisplayOutput).data['text/markdown'], 'foo', 'Incorrect output value');
+        const displayId = (cells[1].outputs[0] as CellDisplayOutput).metadata?.custom?.transient?.display_id;
+        assert.ok(displayId, 'Display id not present in metadata');
+
+        await insertPythonCell(
+            dedent`
+                    dh.update(Markdown("bar"))
+                    print('hello')`,
+            2
+        );
+        await executeActiveDocument();
+        cells = vscodeNotebook.activeNotebookEditor?.document.cells!;
+        await waitForCondition(
+            async () => assertHasExecutionCompletedSuccessfully(cells[2]),
+            15_000,
+            'Cell did not get executed'
+        );
+
+        assert.equal(cells[0].outputs.length, 0, 'Incorrect number of output');
+        assert.equal(cells[1].outputs.length, 1, 'Incorrect number of output');
+        assert.equal(cells[2].outputs.length, 1, 'Incorrect number of output');
+        assert.equal(cells[1].outputs[0].outputKind, vscodeNotebookEnums.CellOutputKind.Rich, 'Incorrect output type');
+        assert.equal((cells[1].outputs[0] as CellDisplayOutput).data['text/markdown'], 'bar', 'Incorrect output value');
+        assertHasTextOutputInVSCode(cells[2], 'hello', 0, false);
+    });
 });
diff --git a/src/test/datascience/notebook/helper.ts b/src/test/datascience/notebook/helper.ts
@@ -90,12 +90,6 @@ export async function insertPythonCell(source: string, index?: number) {
         ])
     );
 }
-export async function insertPythonCellAndWait(source: string, index?: number) {
-    await insertPythonCell(source, index);
-}
-export async function insertMarkdownCellAndWait(source: string) {
-    await insertMarkdownCell(source);
-}
 export async function deleteCell(index: number = 0) {
     const { vscodeNotebook } = await getServices();
     const activeEditor = vscodeNotebook.activeNotebookEditor;
diff --git a/src/test/datascience/notebook/interrupRestart.ds.test.ts b/src/test/datascience/notebook/interrupRestart.ds.test.ts
@@ -22,7 +22,7 @@ import {
     closeNotebooks,
     closeNotebooksAndCleanUpAfterTests,
     executeActiveDocument,
-    insertPythonCellAndWait,
+    insertPythonCell,
     startJupyter,
     trustAllNotebooks,
     waitForTextOutputInVSCode
@@ -82,7 +82,7 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors (slow)',
     });
 
     test('Cancelling token will cancel cell execution', async () => {
-        await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
+        await insertPythonCell('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
         const cell = vscEditor.document.cells[0];
         const appShell = api.serviceContainer.get<IApplicationShell>(IApplicationShell);
         const showInformationMessage = sinon.stub(appShell, 'showInformationMessage');
@@ -116,7 +116,7 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors (slow)',
         }
     });
     test('Restarting kernel will cancel cell execution & we can re-run a cell', async () => {
-        await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
+        await insertPythonCell('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
         const cell = vscEditor.document.cells[0];
 
         await executeActiveDocument();
diff --git a/src/test/datascience/notebook/notebookEditorProvider.ds.test.ts b/src/test/datascience/notebook/notebookEditorProvider.ds.test.ts
@@ -19,7 +19,7 @@ import {
     canRunTests,
     closeNotebooksAndCleanUpAfterTests,
     createTemporaryNotebook,
-    insertMarkdownCellAndWait,
+    insertMarkdownCell,
     trustAllNotebooks
 } from './helper';
 
@@ -147,7 +147,7 @@ suite('DataScience - VSCode Notebook (Editor Provider)', function () {
         const notebookClosed = createEventHandler(editorProvider, 'onDidCloseNotebookEditor', disposables);
 
         const notebookEditor = await editorProvider.open(testIPynb);
-        await insertMarkdownCellAndWait('1'); // Make the file dirty (so it gets pinned).
+        await insertMarkdownCell('1'); // Make the file dirty (so it gets pinned).
         await notebookOpened.assertFired();
         await activeNotebookChanged.assertFired();
         assert.equal(activeNotebookChanged.first, notebookEditor);
@@ -194,7 +194,7 @@ suite('DataScience - VSCode Notebook (Editor Provider)', function () {
         const notebookClosed = createEventHandler(editorProvider, 'onDidCloseNotebookEditor', disposables);
 
         const editor1 = await editorProvider.open(testIPynb);
-        await insertMarkdownCellAndWait('1'); // Make the file dirty (so it gets pinned).
+        await insertMarkdownCell('1'); // Make the file dirty (so it gets pinned).
         await notebookOpened.assertFired();
         await activeNotebookChanged.assertFired();
         assert.equal(activeNotebookChanged.first, editor1);
@@ -203,7 +203,7 @@ suite('DataScience - VSCode Notebook (Editor Provider)', function () {
         // Open another notebook.
         const testIPynb2 = Uri.file(await createTemporaryNotebook(templateIPynb, disposables));
         const editor2 = await editorProvider.open(testIPynb2);
-        await insertMarkdownCellAndWait('1'); // Make the file dirty (so it gets pinned).
+        await insertMarkdownCell('1'); // Make the file dirty (so it gets pinned).
 
         await notebookOpened.assertFiredExactly(2);
         await activeNotebookChanged.assertFiredAtLeast(2);
diff --git a/src/test/datascience/notebook/saving.ds.test.ts b/src/test/datascience/notebook/saving.ds.test.ts
