Update other cells in cell execution

Author: DonJayamanne

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;

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,
@@ -370,12 +374,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)) {

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> {

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>;

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;
     }