Update cell output and metadata using Edit API

Author: DonJayamanne

src/client/datascience/jupyter/kernels/cellExecution.ts

@@ -6,8 +6,9 @@
 import { nbformat } from '@jupyterlab/coreutils';
 import type { KernelMessage } from '@jupyterlab/services/lib/kernel/messages';
 import { CancellationToken, CellOutputKind, CellStreamOutput, NotebookCell, NotebookCellRunState } from 'vscode';
+import type { NotebookEditor as VSCNotebookEditor } from '../../../../../types/vscode-proposed';
 import { concatMultilineString, formatStreamText } from '../../../../datascience-ui/common';
-import { IApplicationShell } from '../../../common/application/types';
+import { IApplicationShell, IVSCodeNotebook } from '../../../common/application/types';
 import { traceInfo, traceWarning } from '../../../common/logger';
 import { RefBool } from '../../../common/refBool';
 import { createDeferred } from '../../../common/utils/async';
@@ -41,12 +42,14 @@ export class CellExecutionFactory {
         private readonly contentProvider: INotebookContentProvider,
         private readonly errorHandler: IDataScienceErrorHandler,
         private readonly editorProvider: INotebookEditorProvider,
-        private readonly appShell: IApplicationShell
+        private readonly appShell: IApplicationShell,
+        private readonly vscNotebook: IVSCodeNotebook
     ) {}
 
     public create(cell: NotebookCell) {
         // tslint:disable-next-line: no-use-before-declare
         return CellExecution.fromCell(
+            this.vscNotebook.notebookEditors.find((e) => e.document === cell.notebook)!,
             cell,
             this.contentProvider,
             this.errorHandler,
@@ -87,6 +90,7 @@ export class CellExecution {
     private _completed?: boolean;
 
     private constructor(
+        public readonly editor: VSCNotebookEditor,
         public readonly cell: NotebookCell,
         private readonly contentProvider: INotebookContentProvider,
         private readonly errorHandler: IDataScienceErrorHandler,
@@ -98,19 +102,20 @@ export class CellExecution {
     }
 
     public static fromCell(
+        editor: VSCNotebookEditor,
         cell: NotebookCell,
         contentProvider: INotebookContentProvider,
         errorHandler: IDataScienceErrorHandler,
         editorProvider: INotebookEditorProvider,
         appService: IApplicationShell
     ) {
-        return new CellExecution(cell, contentProvider, errorHandler, editorProvider, appService);
+        return new CellExecution(editor, cell, contentProvider, errorHandler, editorProvider, appService);
     }
 
     public start(kernelPromise: Promise<IKernel>, notebook: INotebook) {
         this.started = true;
         // Ensure we clear the cell state and trigger a change.
-        clearCellForExecution(this.cell);
+        clearCellForExecution(this.editor, this.cell);
         this.cell.metadata.runStartTime = new Date().getTime();
         this.stopWatch.reset();
         // Changes to metadata must be saved in ipynb, hence mark doc has dirty.
@@ -147,7 +152,7 @@ export class CellExecution {
     private completedWithErrors(error: Partial<Error>) {
         this.sendPerceivedCellExecute();
         this.cell.metadata.lastRunDuration = this.stopWatch.elapsedTime;
-        updateCellWithErrorStatus(this.cell, error);
+        updateCellWithErrorStatus(this.editor, this.cell, error);
         this.contentProvider.notifyChangesToDocument(this.cell.notebook);
         this.errorHandler.handleError((error as unknown) as Error).ignoreErrors();
 
@@ -167,7 +172,7 @@ export class CellExecution {
 
         this.cell.metadata.statusMessage = '';
         this.cell.metadata.lastRunDuration = this.stopWatch.elapsedTime;
-        updateCellExecutionTimes(this.cell, {
+        updateCellExecutionTimes(this.editor, this.cell, {
             startTime: this.cell.metadata.runStartTime,
             duration: this.cell.metadata.lastRunDuration
         });
@@ -351,7 +356,7 @@ export class CellExecution {
 
             // Set execution count, all messages should have it
             if ('execution_count' in msg.content && typeof msg.content.execution_count === 'number') {
-                if (updateCellExecutionCount(this.cell, msg.content.execution_count)) {
+                if (updateCellExecutionCount(this.editor, this.cell, msg.content.execution_count)) {
                     shouldUpdate = true;
                 }
             }
@@ -442,7 +447,7 @@ export class CellExecution {
 
     private handleExecuteInput(msg: KernelMessage.IExecuteInputMsg, _clearState: RefBool) {
         if (msg.content.execution_count) {
-            updateCellExecutionCount(this.cell, msg.content.execution_count);
+            updateCellExecutionCount(this.editor, this.cell, msg.content.execution_count);
         }
     }
 
@@ -518,7 +523,7 @@ export class CellExecution {
 
             // Set execution count, all messages should have it
             if ('execution_count' in msg.content && typeof msg.content.execution_count === 'number') {
-                updateCellExecutionCount(this.cell, msg.content.execution_count);
+                updateCellExecutionCount(this.editor, this.cell, msg.content.execution_count);
             }
 
             // Send this event.

src/client/datascience/jupyter/kernels/kernel.ts

@@ -17,7 +17,7 @@ import {
     Uri
 } from 'vscode';
 import { ServerStatus } from '../../../../datascience-ui/interactive-common/mainState';
-import { IApplicationShell, ICommandManager } from '../../../common/application/types';
+import { IApplicationShell, ICommandManager, IVSCodeNotebook } from '../../../common/application/types';
 import { traceError } from '../../../common/logger';
 import { IDisposableRegistry } from '../../../common/types';
 import { createDeferred, Deferred } from '../../../common/utils/async';
@@ -86,7 +86,8 @@ export class Kernel implements IKernel {
         editorProvider: INotebookEditorProvider,
         private readonly kernelProvider: IKernelProvider,
         private readonly kernelSelectionUsage: IKernelSelectionUsage,
-        appShell: IApplicationShell
+        appShell: IApplicationShell,
+        vscNotebook: IVSCodeNotebook
     ) {
         this.kernelExecution = new KernelExecution(
             kernelProvider,
@@ -96,7 +97,8 @@ export class Kernel implements IKernel {
             contentProvider,
             editorProvider,
             kernelSelectionUsage,
-            appShell
+            appShell,
+            vscNotebook
         );
     }
     public async executeCell(cell: NotebookCell): Promise<void> {

src/client/datascience/jupyter/kernels/kernelExecution.ts

@@ -5,7 +5,7 @@
 
 import { KernelMessage } from '@jupyterlab/services';
 import { NotebookCell, NotebookCellRunState, NotebookDocument } from 'vscode';
-import { IApplicationShell, ICommandManager } from '../../../common/application/types';
+import { IApplicationShell, ICommandManager, IVSCodeNotebook } from '../../../common/application/types';
 import { IDisposable } from '../../../common/types';
 import { noop } from '../../../common/utils/misc';
 import { IInterpreterService } from '../../../interpreter/contracts';
@@ -43,9 +43,16 @@ export class KernelExecution implements IDisposable {
         private readonly contentProvider: INotebookContentProvider,
         editorProvider: INotebookEditorProvider,
         readonly kernelSelectionUsage: IKernelSelectionUsage,
-        readonly appShell: IApplicationShell
+        readonly appShell: IApplicationShell,
+        readonly vscNotebook: IVSCodeNotebook
     ) {
-        this.executionFactory = new CellExecutionFactory(this.contentProvider, errorHandler, editorProvider, appShell);
+        this.executionFactory = new CellExecutionFactory(
+            this.contentProvider,
+            errorHandler,
+            editorProvider,
+            appShell,
+            vscNotebook
+        );
     }
 
     @captureTelemetry(Telemetry.ExecuteNativeCell, undefined, true)
@@ -163,8 +170,11 @@ export class KernelExecution implements IDisposable {
     private onIoPubMessage(document: NotebookDocument, msg: KernelMessage.IIOPubMessage) {
         // tslint:disable-next-line:no-require-imports
         const jupyterLab = require('@jupyterlab/services') as typeof import('@jupyterlab/services');
-        if (jupyterLab.KernelMessage.isUpdateDisplayDataMsg(msg) && handleUpdateDisplayDataMessage(msg, document)) {
-            this.contentProvider.notifyChangesToDocument(document);
+        const editor = this.vscNotebook.notebookEditors.find((e) => e.document === document);
+        if (jupyterLab.KernelMessage.isUpdateDisplayDataMsg(msg) && editor) {
+            if (handleUpdateDisplayDataMessage(msg, editor)) {
+                this.contentProvider.notifyChangesToDocument(document);
+            }
         }
     }
 

src/client/datascience/jupyter/kernels/kernelProvider.ts

@@ -6,7 +6,7 @@
 import * as fastDeepEqual from 'fast-deep-equal';
 import { inject, injectable } from 'inversify';
 import { Uri } from 'vscode';
-import { IApplicationShell, ICommandManager } from '../../../common/application/types';
+import { IApplicationShell, ICommandManager, IVSCodeNotebook } from '../../../common/application/types';
 import { traceInfo, traceWarning } from '../../../common/logger';
 import { IAsyncDisposableRegistry, IConfigurationService, IDisposableRegistry } from '../../../common/types';
 import { IInterpreterService } from '../../../interpreter/contracts';
@@ -30,7 +30,8 @@ export class KernelProvider implements IKernelProvider {
         @inject(INotebookContentProvider) private readonly contentProvider: INotebookContentProvider,
         @inject(INotebookEditorProvider) private readonly editorProvider: INotebookEditorProvider,
         @inject(KernelSelector) private readonly kernelSelectionUsage: IKernelSelectionUsage,
-        @inject(IApplicationShell) private readonly appShell: IApplicationShell
+        @inject(IApplicationShell) private readonly appShell: IApplicationShell,
+        @inject(IVSCodeNotebook) private readonly vscNotebook: IVSCodeNotebook
     ) {}
     public get(uri: Uri): IKernel | undefined {
         return this.kernelsByUri.get(uri.toString())?.kernel;
@@ -57,7 +58,8 @@ export class KernelProvider implements IKernelProvider {
             this.editorProvider,
             this,
             this.kernelSelectionUsage,
-            this.appShell
+            this.appShell,
+            this.vscNotebook
         );
         this.asyncDisposables.push(kernel);
         this.kernelsByUri.set(uri.toString(), { options, kernel });

src/client/datascience/notebook/helpers/executionHelpers.ts

@@ -6,7 +6,7 @@
 import type { nbformat } from '@jupyterlab/coreutils';
 import type { KernelMessage } from '@jupyterlab/services';
 import * as fastDeepEqual from 'fast-deep-equal';
-import { NotebookCell, NotebookCellRunState, NotebookDocument } from 'vscode';
+import type { NotebookCell, NotebookEditor } from '../../../../../types/vscode-proposed';
 import { createErrorOutput } from '../../../../datascience-ui/common/cellFactory';
 import { createIOutputFromCellOutputs, createVSCCellOutputsFromOutputs, translateErrorOutput } from './helpers';
 // tslint:disable-next-line: no-var-requires no-require-imports
@@ -20,8 +20,9 @@ const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed'
  */
 export function handleUpdateDisplayDataMessage(
     msg: KernelMessage.IUpdateDisplayDataMsg,
-    document: NotebookDocument
+    editor: NotebookEditor
 ): boolean {
+    const document = editor.document;
     // Find any cells that have this same display_id
     return (
         document.cells.filter((cellToCheck, index) => {
@@ -56,7 +57,7 @@ export function handleUpdateDisplayDataMessage(
             }
 
             const vscCell = document.cells[index];
-            updateCellOutput(vscCell, changedOutputs);
+            updateCellOutput(editor, vscCell, changedOutputs);
             return true;
         }).length > 0
     );
@@ -65,17 +66,24 @@ export function handleUpdateDisplayDataMessage(
 /**
  * Updates the VSC cell with the error output.
  */
-export function updateCellWithErrorStatus(cell: NotebookCell, ex: Partial<Error>) {
+export function updateCellWithErrorStatus(editor: NotebookEditor, cell: NotebookCell, ex: Partial<Error>) {
+    editor.edit((edit) => {
+        const cellIndex = editor.document.cells.indexOf(cell);
+        edit.replaceMetadata(cellIndex, { ...cell.metadata, runState: vscodeNotebookEnums.NotebookCellRunState.Error });
+        edit.replaceOutput(cellIndex, [translateErrorOutput(createErrorOutput(ex))]);
+    });
     cell.outputs = [translateErrorOutput(createErrorOutput(ex))];
-    cell.metadata.runState = NotebookCellRunState.Error;
 }
 
 /**
  * @returns {boolean} Returns `true` if execution count has changed.
  */
-export function updateCellExecutionCount(vscCell: NotebookCell, executionCount: number): boolean {
-    if (vscCell.metadata.executionOrder !== executionCount && executionCount) {
-        vscCell.metadata.executionOrder = executionCount;
+export function updateCellExecutionCount(editor: NotebookEditor, cell: NotebookCell, executionCount: number): boolean {
+    if (cell.metadata.executionOrder !== executionCount && executionCount) {
+        editor.edit((edit) => {
+            const cellIndex = editor.document.cells.indexOf(cell);
+            edit.replaceMetadata(cellIndex, { ...cell.metadata, executionOrder: executionCount });
+        });
         return true;
     }
     return false;
@@ -87,7 +95,11 @@ export function updateCellExecutionCount(vscCell: NotebookCell, executionCount:
  * Here we update both the VSCode Cell as well as our ICell (cell in our INotebookModel).
  * @returns {(boolean | undefined)} Returns `true` if output has changed.
  */
-export function updateCellOutput(vscCell: NotebookCell, outputs: nbformat.IOutput[]): boolean | undefined {
+export function updateCellOutput(
+    editor: NotebookEditor,
+    vscCell: NotebookCell,
+    outputs: nbformat.IOutput[]
+): boolean | undefined {
     const newOutput = createVSCCellOutputsFromOutputs(outputs);
     // If there was no output and still no output, then nothing to do.
     if (vscCell.outputs.length === 0 && newOutput.length === 0) {
@@ -98,6 +110,6 @@ export function updateCellOutput(vscCell: NotebookCell, outputs: nbformat.IOutpu
     if (vscCell.outputs.length === newOutput.length && fastDeepEqual(vscCell.outputs, newOutput)) {
         return;
     }
-    vscCell.outputs = newOutput;
+    editor.edit((edit) => edit.replaceOutput(0, newOutput));
     return true;
 }