Ability to clear cell output in NativeNotebook (#12307)

For #10496
For #12274

Fixing #12274 only for native notebooks here. Was an easy fix.

Author: DonJayamanne

src/client/common/constants.ts

@@ -111,6 +111,7 @@ export function isUnitTestExecution(): boolean {
 
 // Temporary constant, used to indicate whether we're using custom editor api or not.
 export const UseCustomEditorApi = Symbol('USE_CUSTOM_EDITOR');
+export const UseNativeEditorApi = Symbol('USE_NATIVEEDITOR');
 export const UseProposedApi = Symbol('USE_VSC_PROPOSED_API');
 
 export * from '../constants';

src/client/datascience/interactive-ipynb/nativeEditorStorage.ts

@@ -21,6 +21,7 @@ import { CellState, ICell, IJupyterExecution, IJupyterKernelSpec, INotebookModel
 import detectIndent = require('detect-indent');
 // tslint:disable-next-line:no-require-imports no-var-requires
 import cloneDeep = require('lodash/cloneDeep');
+import { UseNativeEditorApi } from '../../common/constants';
 import { isFileNotFoundError } from '../../common/platform/errors';
 import { sendTelemetryEvent } from '../../telemetry';
 import { pruneCell } from '../common';
@@ -106,6 +107,7 @@ export class NativeEditorNotebookModel implements INotebookModel {
     private _id = uuid();
 
     constructor(
+        private readonly useNativeEditorApi: boolean,
         file: Uri,
         cells: ICell[],
         json: Partial<nbformat.INotebookContent> = {},
@@ -131,6 +133,7 @@ export class NativeEditorNotebookModel implements INotebookModel {
     }
     public clone(file: Uri) {
         return new NativeEditorNotebookModel(
+            this.useNativeEditorApi,
             file,
             cloneDeep(this._state.cells),
             cloneDeep(this._state.notebookJson),
@@ -213,11 +216,19 @@ export class NativeEditorNotebookModel implements INotebookModel {
                 break;
             case 'save':
                 this._state.saveChangeCount = this._state.changeCount;
+                // Trigger event.
+                if (this.useNativeEditorApi) {
+                    changed = true;
+                }
                 break;
             case 'saveAs':
                 this._state.saveChangeCount = this._state.changeCount;
                 this._state.changeCount = this._state.saveChangeCount = 0;
                 this._state.file = change.target;
+                // Trigger event.
+                if (this.useNativeEditorApi) {
+                    changed = true;
+                }
                 break;
             default:
                 break;
@@ -361,6 +372,11 @@ export class NativeEditorNotebookModel implements INotebookModel {
     }
 
     private clearOutputs(): boolean {
+        if (this.useNativeEditorApi) {
+            // Do not create new cells when using native editor.
+            // We'll update the cells in place (cuz undo/redo is handled by VS Code).
+            return true;
+        }
         const newCells = this.cells.map((c) =>
             this.asCell({ ...c, data: { ...c.data, execution_count: null, outputs: [] } })
         );
@@ -496,7 +512,8 @@ export class NativeEditorStorage implements INotebookStorage {
         @inject(ICryptoUtils) private crypto: ICryptoUtils,
         @inject(IExtensionContext) private context: IExtensionContext,
         @inject(IMemento) @named(GLOBAL_MEMENTO) private globalStorage: Memento,
-        @inject(IMemento) @named(WORKSPACE_MEMENTO) private localStorage: Memento
+        @inject(IMemento) @named(WORKSPACE_MEMENTO) private localStorage: Memento,
+        @inject(UseNativeEditorApi) private readonly useNativeEditorApi: boolean
     ) {}
     private static isUntitledFile(file: Uri) {
         return isUntitledFile(file);
@@ -669,7 +686,7 @@ export class NativeEditorStorage implements INotebookStorage {
         } catch (ex) {
             // May not exist at this time. Should always have a single cell though
             traceError(`Failed to load notebook file ${file.toString()}`, ex);
-            return new NativeEditorNotebookModel(file, []);
+            return new NativeEditorNotebookModel(this.useNativeEditorApi, file, []);
         }
     }
 
@@ -720,7 +737,15 @@ export class NativeEditorStorage implements INotebookStorage {
             remapped.splice(0, 0, this.createEmptyCell(uuid()));
         }
         const pythonNumber = json ? await this.extractPythonMainVersion(json) : 3;
-        return new NativeEditorNotebookModel(file, remapped, json, indentAmount, pythonNumber, isInitiallyDirty);
+        return new NativeEditorNotebookModel(
+            this.useNativeEditorApi,
+            file,
+            remapped,
+            json,
+            indentAmount,
+            pythonNumber,
+            isInitiallyDirty
+        );
     }
 
     private getStorageKey(file: Uri): string {

src/client/datascience/notebook/contentProvider.ts

@@ -44,12 +44,6 @@ export class NotebookContentProvider implements VSCodeNotebookContentProvider {
         }
         if (model.isUntitled) {
             await this.commandManager.executeCommand('workbench.action.files.saveAs', document.uri);
-            // tslint:disable-next-line: no-suspicious-comment
-            //TODO: VSC doesn't handle this well.
-            // What if user doesn't save it.
-            if (model.isUntitled) {
-                throw new Error('Not saved');
-            }
         } else {
             await this.notebookStorage.save(model, cancellation);
         }

src/client/datascience/notebook/executionService.ts

@@ -11,6 +11,7 @@ import type { NotebookCell, NotebookDocument } from 'vscode-proposed';
 import { ICommandManager } from '../../common/application/types';
 import { wrapCancellationTokens } from '../../common/cancellation';
 import '../../common/extensions';
+import { IDisposable } from '../../common/types';
 import { createDeferred } from '../../common/utils/async';
 import { noop } from '../../common/utils/misc';
 import { StopWatch } from '../../common/utils/stopWatch';
@@ -152,6 +153,7 @@ export class NotebookExecutionService implements INotebookExecutionService {
         cell.metadata.runState = vscodeNotebookEnums.NotebookCellRunState.Running;
 
         let subscription: Subscription | undefined;
+        let modelClearedEventHandler: IDisposable | undefined;
         try {
             nb.clear(cell.uri.fsPath); // NOSONAR
             const observable = nb.executeObservable(
@@ -163,6 +165,15 @@ export class NotebookExecutionService implements INotebookExecutionService {
             );
             subscription = observable?.subscribe(
                 (cells) => {
+                    if (!modelClearedEventHandler) {
+                        modelClearedEventHandler = model.changed((e) => {
+                            if (e.kind === 'clear') {
+                                // If cell output has been cleared, then clear the output in the observed executable cell.
+                                // Else if user clears output while execuitng a cell, we add it back.
+                                cells.forEach((c) => (c.data.outputs = []));
+                            }
+                        });
+                    }
                     const rawCellOutput = cells
                         .filter((item) => item.id === cell.uri.fsPath)
                         .flatMap((item) => (item.data.outputs as unknown) as nbformat.IOutput[])
@@ -198,7 +209,6 @@ export class NotebookExecutionService implements INotebookExecutionService {
                     const notebookCellModel = findMappedNotebookCellModel(document, cell, model.cells);
                     updateCellExecutionTimes(
                         notebookCellModel,
-                        model,
                         cell.metadata.runStartTime,
                         cell.metadata.lastRunDuration
                     );
@@ -221,6 +231,7 @@ export class NotebookExecutionService implements INotebookExecutionService {
             updateCellWithErrorStatus(cell, ex);
             this.errorHandler.handleError(ex).ignoreErrors();
         } finally {
+            modelClearedEventHandler?.dispose(); // NOSONAR
             subscription?.unsubscribe(); // NOSONAR
             // Ensure we remove the cancellation.
             const cancellations = this.pendingExecutionCancellations.get(document.uri.fsPath);

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

@@ -18,9 +18,13 @@ import {
 } from '../../../common/application/types';
 import { MARKDOWN_LANGUAGE, PYTHON_LANGUAGE } from '../../../common/constants';
 import { traceError } from '../../../logging';
-import { ICell, INotebookModel } from '../../types';
+import { INotebookModel } from '../../types';
 import { findMappedNotebookCellModel } from './cellMappers';
-import { createCellFromVSCNotebookCell, createVSCCellOutputsFromOutputs } from './helpers';
+import {
+    createCellFromVSCNotebookCell,
+    createVSCCellOutputsFromOutputs,
+    updateVSCNotebookCellMetadata
+} from './helpers';
 
 /**
  * If a VS Code cell changes, then ensure we update the corresponding cell in our INotebookModel.
@@ -32,9 +36,7 @@ export function updateCellModelWithChangesToVSCCell(
 ) {
     switch (change.type) {
         case 'changeCellOutputs':
-            // We're not interested in changes to cell output as this happens as a result of us pushing changes to the notebook.
-            // I.e. cell output is already in our INotebookModel.
-            return;
+            return clearCellOutput(change, model);
         case 'changeCellLanguage':
             return changeCellLanguage(change, model);
         case 'changeCells':
@@ -45,6 +47,31 @@ export function updateCellModelWithChangesToVSCCell(
     }
 }
 
+/**
+ * We're not interested in changes to cell output as this happens as a result of us pushing changes to the notebook.
+ * I.e. cell output is already in our INotebookModel.
+ * However we are interested in cell output being cleared (when user clears output).
+ */
+function clearCellOutput(change: NotebookCellOutputsChangeEvent, model: INotebookModel) {
+    if (!change.cells.every((cell) => cell.outputs.length === 0)) {
+        return;
+    }
+
+    // If a cell has been cleared, then clear the corresponding ICell (cell in INotebookModel).
+    change.cells.forEach((vscCell) => {
+        const cell = findMappedNotebookCellModel(change.document, vscCell, model.cells);
+        cell.data.outputs = [];
+        updateVSCNotebookCellMetadata(vscCell.metadata, cell);
+        model.update({
+            source: 'user',
+            kind: 'clear',
+            oldDirty: model.isDirty,
+            newDirty: true,
+            oldCells: [cell]
+        });
+    });
+}
+
 function changeCellLanguage(change: NotebookCellLanguageChangeEvent, model: INotebookModel) {
     const cellModel = findMappedNotebookCellModel(change.document, change.cell, model.cells);
 
@@ -57,26 +84,16 @@ function changeCellLanguage(change: NotebookCellLanguageChangeEvent, model: INot
         return;
     }
 
-    const newCellData = createCellFrom(cellModel.data, change.language === MARKDOWN_LANGUAGE ? 'markdown' : 'code');
-    const newCell: ICell = {
-        ...cellModel,
-        data: newCellData
-    };
-
-    model.update({
-        source: 'user',
-        kind: 'modify',
-        newCells: [newCell],
-        oldCells: [cellModel],
-        oldDirty: model.isDirty,
-        newDirty: true
-    });
-
+    const cellData = createCellFrom(cellModel.data, change.language === MARKDOWN_LANGUAGE ? 'markdown' : 'code');
     // tslint:disable-next-line: no-any
-    change.cell.outputs = createVSCCellOutputsFromOutputs(newCellData.outputs as any);
+    change.cell.outputs = createVSCCellOutputsFromOutputs(cellData.outputs as any);
     change.cell.metadata.executionOrder = undefined;
     change.cell.metadata.hasExecutionOrder = change.language !== MARKDOWN_LANGUAGE; // Do not check for Python, to support other languages
     change.cell.metadata.runnable = change.language !== MARKDOWN_LANGUAGE; // Do not check for Python, to support other languages
+
+    // Create a new cell & replace old one.
+    const oldCellIndex = model.cells.indexOf(cellModel);
+    model.cells[oldCellIndex] = createCellFromVSCNotebookCell(change.document, change.cell, model);
 }
 
 function handleChangesToCells(change: NotebookCellsChangeEvent, model: INotebookModel) {
@@ -135,42 +152,22 @@ function handleCellMove(change: NotebookCellsChangeEvent, model: INotebookModel)
     const cellToSwap = findMappedNotebookCellModel(change.document, insertChange.items[0]!, model.cells);
     const cellToSwapWith = model.cells[insertChange.start];
     assert.notEqual(cellToSwap, cellToSwapWith, 'Cannot swap cell with the same cell');
-    model.update({
-        source: 'user',
-        kind: 'swap',
-        oldDirty: model.isDirty,
-        newDirty: true,
-        firstCellId: cellToSwap.id,
-        secondCellId: cellToSwapWith.id
-    });
+
+    const indexOfCellToSwap = model.cells.indexOf(cellToSwap);
+    model.cells[insertChange.start] = cellToSwap;
+    model.cells[indexOfCellToSwap] = cellToSwapWith;
 }
 function handleCellInsertion(change: NotebookCellsChangeEvent, model: INotebookModel) {
     assert.equal(change.changes.length, 1, 'When inserting cells we must have only 1 change');
     assert.equal(change.changes[0].items.length, 1, 'Insertion of more than 1 cell is not supported');
     const insertChange = change.changes[0];
     const cell = change.changes[0].items[0];
     const newCell = createCellFromVSCNotebookCell(change.document, cell, model);
-
-    model.update({
-        source: 'user',
-        kind: 'insert',
-        newDirty: true,
-        oldDirty: model.isDirty,
-        index: insertChange.start,
-        cell: newCell
-    });
+    model.cells.splice(insertChange.start, 0, newCell);
 }
 function handleCellDelete(change: NotebookCellsChangeEvent, model: INotebookModel) {
     assert.equal(change.changes.length, 1, 'When deleting cells we must have only 1 change');
     const deletionChange = change.changes[0];
     assert.equal(deletionChange.deletedCount, 1, 'Deleting more than one cell is not supported');
-    const cellToDelete = model.cells[deletionChange.start];
-    model.update({
-        source: 'user',
-        kind: 'remove',
-        oldDirty: model.isDirty,
-        newDirty: true,
-        cell: cellToDelete,
-        index: deletionChange.start
-    });
+    model.cells.splice(deletionChange.start, 1);
 }

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

@@ -8,7 +8,6 @@ import type { KernelMessage } from '@jupyterlab/services';
 import { NotebookCell, NotebookCellRunState, NotebookDocument } from 'vscode';
 import { IBaseCellVSCodeMetadata } from '../../../../../types/@jupyterlab_coreutils_nbformat';
 import { createErrorOutput } from '../../../../datascience-ui/common/cellFactory';
-import { INotebookModelModifyChange } from '../../interactive-common/interactiveWindowTypes';
 import { ICell, INotebookModel } from '../../types';
 import { findMappedNotebookCell } from './cellMappers';
 import { createVSCCellOutputsFromOutputs, translateErrorOutput, updateVSCNotebookCellMetadata } from './helpers';
@@ -118,51 +117,23 @@ export function updateCellOutput(vscCell: NotebookCell, cell: ICell, outputs: nb
 /**
  * Store execution start and end times in ISO format for portability.
  */
-export function updateCellExecutionTimes(
-    notebookCellModel: ICell,
-    model: INotebookModel,
-    startTime?: number,
-    duration?: number
-) {
+export function updateCellExecutionTimes(notebookCellModel: ICell, startTime?: number, duration?: number) {
     const startTimeISO = startTime ? new Date(startTime).toISOString() : undefined;
     const endTimeISO = duration && startTime ? new Date(startTime + duration).toISOString() : undefined;
-    updateCellMetadata(
-        notebookCellModel,
-        {
-            end_execution_time: endTimeISO,
-            start_execution_time: startTimeISO
-        },
-        model
-    );
+    updateCellMetadata(notebookCellModel, {
+        end_execution_time: endTimeISO,
+        start_execution_time: startTimeISO
+    });
 }
 
-export function updateCellMetadata(
-    notebookCellModel: ICell,
-    metadata: Partial<IBaseCellVSCodeMetadata>,
-    model: INotebookModel
-) {
+export function updateCellMetadata(notebookCellModel: ICell, metadata: Partial<IBaseCellVSCodeMetadata>) {
     const originalVscodeMetadata: IBaseCellVSCodeMetadata = notebookCellModel.data.metadata.vscode || {};
     // Update our model with the new metadata stored in jupyter.
-    const newCell: ICell = {
-        ...notebookCellModel,
-        data: {
-            ...notebookCellModel.data,
-            metadata: {
-                ...notebookCellModel.data.metadata,
-                vscode: {
-                    ...originalVscodeMetadata,
-                    ...metadata
-                }
-            }
+    notebookCellModel.data.metadata = {
+        ...notebookCellModel.data.metadata,
+        vscode: {
+            ...originalVscodeMetadata,
+            ...metadata
         }
     };
-    const updateCell: INotebookModelModifyChange = {
-        kind: 'modify',
-        newCells: [newCell],
-        oldCells: [notebookCellModel],
-        newDirty: true,
-        oldDirty: model.isDirty === true,
-        source: 'user'
-    };
-    model.update(updateCell);
 }

src/client/datascience/serviceRegistry.ts

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 'use strict';
 import { IExtensionSingleActivationService } from '../activation/types';
-import { UseCustomEditorApi } from '../common/constants';
+import { UseCustomEditorApi, UseNativeEditorApi } from '../common/constants';
 import { NotebookEditorSupport } from '../common/experiments/groups';
 import { StartPage } from '../common/startPage/startPage';
 import { IStartPage } from '../common/startPage/types';
@@ -164,6 +164,7 @@ export function registerTypes(serviceManager: IServiceManager) {
     const inCustomEditorApiExperiment = experiments.inExperiment(NotebookEditorSupport.customEditorExperiment);
     const usingCustomEditor = inCustomEditorApiExperiment;
     serviceManager.addSingletonInstance<boolean>(UseCustomEditorApi, usingCustomEditor);
+    serviceManager.addSingletonInstance<boolean>(UseNativeEditorApi, useVSCodeNotebookAPI);
 
     // This condition is temporary.
     const notebookEditorProvider = useVSCodeNotebookAPI ? NotebookEditorProvider : usingCustomEditor ? NativeEditorProvider : NativeEditorProviderOld;