Update cell output and metadata using Edit API (#13737)

Author: DonJayamanne

.vscode/launch.json

@@ -143,7 +143,7 @@
                 "--extensionTestsPath=${workspaceFolder}/out/test"
             ],
             "env": {
-                "VSC_PYTHON_CI_TEST_GREP": "Language Server:" 
+                "VSC_PYTHON_CI_TEST_GREP": "Language Server:"
             },
             "stopOnEntry": false,
             "sourceMaps": true,
@@ -166,6 +166,7 @@
             "env": {
                 "VSC_PYTHON_CI_TEST_GREP": "", // Modify this to run a subset of the single workspace tests
                 "VSC_PYTHON_CI_TEST_INVERT_GREP": "", // Initialize this to invert the grep (exclude tests with value defined in grep).
+                "CI_PYTHON_PATH": "<PythonPath>", // Initialize this to invert the grep (exclude tests with value defined in grep).
 
                 "VSC_PYTHON_LOAD_EXPERIMENTS_FROM_FILE": "true",
                 "TEST_FILES_SUFFIX": "ds.test"

src/client/common/utils/decorators.ts

@@ -178,11 +178,11 @@ export function cache(expiryDurationMs: number) {
  * @param {string} [scopeName] Scope for the error message to be logged along with the error.
  * @returns void
  */
-export function swallowExceptions(scopeName: string) {
+export function swallowExceptions(scopeName?: string) {
     // tslint:disable-next-line:no-any no-function-expression
     return function (_target: any, propertyName: string, descriptor: TypedPropertyDescriptor<any>) {
         const originalMethod = descriptor.value!;
-        const errorMessage = `Python Extension (Error in ${scopeName}, method:${propertyName}):`;
+        const errorMessage = `Python Extension (Error in ${scopeName || propertyName}, method:${propertyName}):`;
         // tslint:disable-next-line:no-any no-function-expression
         descriptor.value = function (...args: any[]) {
             try {

src/client/datascience/jupyter/kernels/cellExecution.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> {
@@ -107,11 +109,11 @@ export class Kernel implements IKernel {
         await this.start({ disableUI: false, token: this.startCancellation.token });
         await this.kernelExecution.executeAllCells(document);
     }
-    public cancelCell(cell: NotebookCell) {
+    public async cancelCell(cell: NotebookCell) {
         this.startCancellation.cancel();
-        this.kernelExecution.cancelCell(cell);
+        await this.kernelExecution.cancelCell(cell);
     }
-    public cancelAllCells(document: NotebookDocument) {
+    public async cancelAllCells(document: NotebookDocument) {
         this.startCancellation.cancel();
         this.kernelExecution.cancelAllCells(document);
     }

src/client/datascience/jupyter/kernels/kernel.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)
@@ -78,11 +85,17 @@ export class KernelExecution implements IDisposable {
         if (this.documentExecutions.has(document)) {
             return;
         }
+        const editor = this.vscNotebook.notebookEditors.find((item) => item.document === document);
+        if (!editor) {
+            return;
+        }
         const cancelTokenSource = new MultiCancellationTokenSource();
         this.documentExecutions.set(document, cancelTokenSource);
         const kernel = this.getKernel(document);
-        document.metadata.runState = vscodeNotebookEnums.NotebookRunState.Running;
 
+        await editor.edit((edit) =>
+            edit.replaceMetadata({ ...document.metadata, runState: vscodeNotebookEnums.NotebookRunState.Running })
+        );
         const codeCellsToExecute = document.cells
             .filter((cell) => cell.cellKind === vscodeNotebookEnums.CellKind.Code)
             .filter((cell) => cell.document.getText().trim().length > 0)
@@ -98,35 +111,31 @@ export class KernelExecution implements IDisposable {
         );
 
         try {
-            let executingAPreviousCellHasFailed = false;
-            await codeCellsToExecute.reduce(
-                (previousPromise, cellToExecute) =>
-                    previousPromise.then((previousCellState) => {
-                        // If a previous cell has failed or execution cancelled, the get out.
-                        if (
-                            executingAPreviousCellHasFailed ||
-                            cancelTokenSource.token.isCancellationRequested ||
-                            previousCellState === vscodeNotebookEnums.NotebookCellRunState.Error
-                        ) {
-                            executingAPreviousCellHasFailed = true;
-                            codeCellsToExecute.forEach((cell) => cell.cancel()); // Cancel pending cells.
-                            return;
-                        }
-                        const result = this.executeIndividualCell(kernel, cellToExecute);
-                        result.finally(() => this.cellExecutions.delete(cellToExecute.cell)).catch(noop);
-                        return result;
-                    }),
-                Promise.resolve<NotebookCellRunState | undefined>(undefined)
-            );
+            for (const cellToExecute of codeCellsToExecute) {
+                const result = this.executeIndividualCell(kernel, cellToExecute);
+                result.finally(() => this.cellExecutions.delete(cellToExecute.cell)).catch(noop);
+                const executionResult = await result;
+                // If a cell has failed or execution cancelled, the get out.
+                if (
+                    cancelTokenSource.token.isCancellationRequested ||
+                    executionResult === vscodeNotebookEnums.NotebookCellRunState.Error
+                ) {
+                    await Promise.all(codeCellsToExecute.map((cell) => cell.cancel())); // Cancel pending cells.
+                    break;
+                }
+            }
         } finally {
+            await Promise.all(codeCellsToExecute.map((cell) => cell.cancel())); // Cancel pending cells.
             this.documentExecutions.delete(document);
-            document.metadata.runState = vscodeNotebookEnums.NotebookRunState.Idle;
+            await editor.edit((edit) =>
+                edit.replaceMetadata({ ...document.metadata, runState: vscodeNotebookEnums.NotebookRunState.Idle })
+            );
         }
     }
 
-    public cancelCell(cell: NotebookCell): void {
+    public async cancelCell(cell: NotebookCell) {
         if (this.cellExecutions.get(cell)) {
-            this.cellExecutions.get(cell)!.cancel();
+            await this.cellExecutions.get(cell)!.cancel();
         }
     }
 
@@ -160,11 +169,14 @@ export class KernelExecution implements IDisposable {
         return kernel;
     }
 
-    private onIoPubMessage(document: NotebookDocument, msg: KernelMessage.IIOPubMessage) {
+    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');
-        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 (await handleUpdateDisplayDataMessage(msg, editor)) {
+                this.contentProvider.notifyChangesToDocument(document);
+            }
         }
     }
 
@@ -195,7 +207,7 @@ export class KernelExecution implements IDisposable {
         );
 
         // Start execution
-        cellExecution.start(kernelPromise, this.notebook);
+        await cellExecution.start(kernelPromise, this.notebook);
 
         // The result promise will resolve when complete.
         try {

src/client/datascience/jupyter/kernels/kernelExecution.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/jupyter/kernels/kernelProvider.ts

@@ -45,8 +45,8 @@ export class NotebookContentProvider implements INotebookContentProvider {
         @inject(NotebookEditorCompatibilitySupport)
         private readonly compatibilitySupport: NotebookEditorCompatibilitySupport
     ) {}
-    public notifyChangesToDocument(document: NotebookDocument) {
-        this.notebookChanged.fire({ document });
+    public notifyChangesToDocument(_document: NotebookDocument) {
+        // this.notebookChanged.fire({ document });
     }
     public async resolveNotebook(_document: NotebookDocument, _webview: NotebookCommunication): Promise<void> {
         // Later