Sync VSCode output Cell Output & add INotebookEditor/INotebookEditorProvider

src/client/datascience/notebook/executionHelpers.ts

@@ -7,9 +7,10 @@ import type { nbformat } from '@jupyterlab/coreutils';
 import type { KernelMessage } from '@jupyterlab/services';
 import { NotebookCell, NotebookCellRunState, NotebookDocument } from 'vscode';
 import { createErrorOutput } from '../../../datascience-ui/common/cellFactory';
+import { IDisposable } from '../../common/types';
 import { INotebookModelModifyChange } from '../interactive-common/interactiveWindowTypes';
 import { ICell, INotebookModel } from '../types';
-import { cellOutputsToVSCCellOutputs, translateErrorOutput } from './helpers';
+import { cellOutputsToVSCCellOutputs, findMappedNotebookCellData, translateErrorOutput } from './helpers';
 
 export function hasTransientOutputForAnotherCell(output?: nbformat.IOutput) {
     return (
@@ -22,6 +23,7 @@ export function hasTransientOutputForAnotherCell(output?: nbformat.IOutput) {
         Object.keys((output as any).transient).length > 0
     );
 }
+
 /**
  * Updates the cell in notebook model as well as the notebook document.
  * Update notebook document so UI is updated accordingly.
@@ -76,6 +78,47 @@ export function updateCellWithErrorStatus(cell: NotebookCell, ex: Partial<Error>
     cell.metadata.runState = NotebookCellRunState.Error;
 }
 
+/**
+ * Responsible for syncing changes from our model into the VS Code cells.
+ * Eg. when executing a cell, we update our model with the output, and here we react to those events and update the VS Code output.
+ * This way, all updates to VSCode cells can happen in one place (here), and we can focus on updating just the Cell model with the data.
+ * Here we only keep the outputs in sync. The assumption is that we won't be adding cells directly.
+ * If adding cells and the like then please use VSC api to manipulate cells, else we have 2 ways of doing the same thing and that could lead to issues.
+ */
+export function monitorModelCellOutputChangesAndUpdateNotebookDocument(
+    document: NotebookDocument,
+    model: INotebookModel
+): IDisposable {
+    return model.changed((change) => {
+        // We're only interested in updates to cells.
+        if (change.kind !== 'modify') {
+            return;
+        }
+        for (const cell of change.newCells) {
+            const uiCellToUpdate = findMappedNotebookCellData(cell, document.cells);
+            if (!uiCellToUpdate) {
+                continue;
+            }
+            const newOutput = Array.isArray(cell.data.outputs)
+                ? // tslint:disable-next-line: no-any
+                  cellOutputsToVSCCellOutputs(cell.data.outputs as any)
+                : [];
+            // If there were no cells and still no cells, nothing to update.
+            if (newOutput.length === 0 && uiCellToUpdate.outputs.length === 0) {
+                return;
+            }
+            // If no changes in output, then nothing to do.
+            if (
+                newOutput.length === uiCellToUpdate.outputs.length &&
+                JSON.stringify(newOutput) === JSON.stringify(uiCellToUpdate.outputs)
+            ) {
+                return;
+            }
+            uiCellToUpdate.outputs = newOutput;
+        }
+    });
+}
+
 /**
  * Updates our Cell Model with the cell output.
  * As we execute a cell we get output from jupyter. This code will ensure the cell is updated with the output.

src/client/datascience/notebook/notebookEditor.ts

@@ -0,0 +1,93 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { Event, EventEmitter, notebook, Uri, WebviewPanel } from 'vscode';
+import { INotebook, INotebookEditor, INotebookModel } from '../types';
+
+export class NotebookEditor implements INotebookEditor {
+    public get onDidChangeViewState(): Event<void> {
+        return this.changedViewState.event;
+    }
+    public get closed(): Event<INotebookEditor> {
+        return this._closed.event;
+    }
+    public get modified(): Event<INotebookEditor> {
+        return this._modified.event;
+    }
+
+    public get executed(): Event<INotebookEditor> {
+        return this._executed.event;
+    }
+    public get saved(): Event<INotebookEditor> {
+        return this._saved.event;
+    }
+    public get isUntitled(): boolean {
+        return this.model.isUntitled;
+    }
+    public get isDirty(): boolean {
+        return this.model.isDirty;
+    }
+    public get file(): Uri {
+        return this.model.file;
+    }
+    public get visible(): boolean {
+        return !this.model.isDisposed;
+    }
+    public get active(): boolean {
+        return notebook.activeNotebookEditor?.document.uri.toString() === this.model.file.toString();
+    }
+    public get onExecutedCode(): Event<string> {
+        return this.executedCode.event;
+    }
+    public notebook?: INotebook | undefined;
+    private changedViewState = new EventEmitter<void>();
+    private _closed = new EventEmitter<INotebookEditor>();
+    private _saved = new EventEmitter<INotebookEditor>();
+    private _executed = new EventEmitter<INotebookEditor>();
+    private _modified = new EventEmitter<INotebookEditor>();
+    private executedCode = new EventEmitter<string>();
+    constructor(public readonly model: INotebookModel) {
+        model.onDidEdit(() => this._modified.fire(this));
+    }
+    public async load(_storage: INotebookModel, _webViewPanel?: WebviewPanel): Promise<void> {
+        // Not used.
+    }
+    public runAllCells(): void {
+        throw new Error('Method not implemented.');
+    }
+    public runSelectedCell(): void {
+        throw new Error('Method not implemented.');
+    }
+    public addCellBelow(): void {
+        throw new Error('Method not implemented.');
+    }
+    public show(): Promise<void> {
+        throw new Error('Method not implemented.');
+    }
+    public startProgress(): void {
+        throw new Error('Method not implemented.');
+    }
+    public stopProgress(): void {
+        throw new Error('Method not implemented.');
+    }
+    public undoCells(): void {
+        throw new Error('Method not implemented.');
+    }
+    public redoCells(): void {
+        throw new Error('Method not implemented.');
+    }
+    public removeAllCells(): void {
+        throw new Error('Method not implemented.');
+    }
+    public interruptKernel(): Promise<void> {
+        throw new Error('Method not implemented.');
+    }
+    public restartKernel(): Promise<void> {
+        throw new Error('Method not implemented.');
+    }
+    public dispose() {
+        // Not required.
+    }
+}

src/client/datascience/notebook/notebookEditorProvider.ts

@@ -0,0 +1,184 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+import { inject, injectable } from 'inversify';
+import { Event, EventEmitter, notebook, NotebookDocument, Uri } from 'vscode';
+import { IExtensionSingleActivationService } from '../../activation/types';
+import { ICommandManager, IWorkspaceService } from '../../common/application/types';
+import '../../common/extensions';
+import { IDisposableRegistry } from '../../common/types';
+import { createDeferred, Deferred } from '../../common/utils/async';
+import { isUri } from '../../common/utils/misc';
+import { sendTelemetryEvent } from '../../telemetry';
+import { Telemetry } from '../constants';
+import { INotebookStorageProvider } from '../interactive-ipynb/notebookStorageProvider';
+import { INotebookEditor, INotebookEditorProvider } from '../types';
+import { monitorModelCellOutputChangesAndUpdateNotebookDocument } from './executionHelpers';
+import { NotebookEditor } from './notebookEditor';
+
+/**
+ * Class responsbile for activating an registering the necessary event handlers in NotebookEditorProvider.
+ */
+@injectable()
+export class NotebookEditorProviderActivation implements IExtensionSingleActivationService {
+    constructor(@inject(INotebookEditorProvider) private readonly provider: INotebookEditorProvider) {}
+    public async activate(): Promise<void> {
+        // The whole purpose is to ensure the NotebookEditorProvider class activates as soon as extension loads.
+        // tslint:disable-next-line: no-use-before-declare
+        if (this.provider instanceof NotebookEditorProvider) {
+            this.provider.activate();
+        }
+    }
+}
+
+/**
+ * Notebook Editor provider used by other parts of DS code.
+ * This is an adapter, that takes the VSCode api for editors (did notebook editors open, close save, etc) and
+ * then exposes them in a manner we expect - i.e. INotebookEditorProvider.
+ * This is also responsible for tracking all notebooks that open and then keeping the VS Code notebook models updated with changes we made to our underlying model.
+ * E.g. when cells are executed the results in our model is updated, this tracks those changes and syncs VSC cells with those updates.
+ */
+@injectable()
+export class NotebookEditorProvider implements INotebookEditorProvider {
+    public get onDidChangeActiveNotebookEditor(): Event<INotebookEditor | undefined> {
+        return this._onDidChangeActiveNotebookEditor.event;
+    }
+    public get onDidCloseNotebookEditor(): Event<INotebookEditor> {
+        return this._onDidCloseNotebookEditor.event;
+    }
+    public get onDidOpenNotebookEditor(): Event<INotebookEditor> {
+        return this._onDidOpenNotebookEditor.event;
+    }
+    public get activeEditor(): INotebookEditor | undefined {
+        return this.editors.find((e) => e.visible && e.active);
+    }
+    public get editors(): INotebookEditor[] {
+        return [...this.openedEditors];
+    }
+    // Note, this constant has to match the value used in the package.json to register the webview custom editor.
+    public static readonly customEditorViewType = 'NativeEditorProvider.ipynb';
+    protected readonly _onDidChangeActiveNotebookEditor = new EventEmitter<INotebookEditor | undefined>();
+    protected readonly _onDidOpenNotebookEditor = new EventEmitter<INotebookEditor>();
+    private readonly _onDidCloseNotebookEditor = new EventEmitter<INotebookEditor>();
+    private openedEditors: Set<INotebookEditor> = new Set<INotebookEditor>();
+    private executedEditors: Set<string> = new Set<string>();
+    private notebookCount: number = 0;
+    private openedNotebookCount: number = 0;
+    private readonly notebookEditors = new Map<NotebookDocument, INotebookEditor>();
+    private readonly notebookEditorsByUri = new Map<string, INotebookEditor>();
+    private readonly notebooksWaitingToBeOpenedByUri = new Map<string, Deferred<INotebookEditor>>();
+    constructor(
+        @inject(INotebookStorageProvider) private readonly storage: INotebookStorageProvider,
+        @inject(IWorkspaceService) private readonly workspace: IWorkspaceService,
+        @inject(ICommandManager) private readonly commandManager: ICommandManager,
+        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry
+    ) {
+        disposables.push(this);
+    }
+    public activate() {
+        this.disposables.push(notebook.onDidOpenNotebookDocument(this.onDidOpenNotebookDocument, this));
+        this.disposables.push(notebook.onDidCloseNotebookDocument(this.onDidCloseNotebookDocument, this));
+
+        // Look through the file system for ipynb files to see how many we have in the workspace. Don't wait
+        // on this though.
+        const findFilesPromise = this.workspace.findFiles('**/*.ipynb');
+        if (findFilesPromise && findFilesPromise.then) {
+            findFilesPromise.then((r) => (this.notebookCount += r.length));
+        }
+    }
+    public dispose() {
+        // Send a bunch of telemetry
+        if (this.openedNotebookCount) {
+            sendTelemetryEvent(Telemetry.NotebookOpenCount, undefined, { count: this.openedNotebookCount });
+        }
+        if (this.executedEditors.size) {
+            sendTelemetryEvent(Telemetry.NotebookRunCount, undefined, { count: this.executedEditors.size });
+        }
+        if (this.notebookCount) {
+            sendTelemetryEvent(Telemetry.NotebookWorkspaceCount, undefined, { count: this.notebookCount });
+        }
+    }
+
+    public async open(file: Uri): Promise<INotebookEditor> {
+        // Wait for editor to get opened up, vscode will notify when it is opened.
+        // Further below.
+        const deferred = createDeferred<INotebookEditor>();
+        this.notebooksWaitingToBeOpenedByUri.set(file.toString(), deferred);
+        deferred.promise.then(() => this.notebooksWaitingToBeOpenedByUri.delete(file.toString())).ignoreErrors();
+        await this.commandManager.executeCommand('vscode.open', file);
+        return deferred.promise;
+    }
+    public async show(_file: Uri): Promise<INotebookEditor | undefined> {
+        // We do not need this.
+        throw new Error('Not supported');
+    }
+    public async createNew(_contents?: string): Promise<INotebookEditor> {
+        // tslint:disable-next-line: no-suspicious-comment
+        // TODO: In another branch.
+        // const model = await this.storage.createNew(contents);
+        // await this.onDidOpenNotebookDocument(model.file);
+        // tslint:disable-next-line: no-suspicious-comment
+        // TODO: Need to do this.
+        // Update number of notebooks in the workspace
+        // this.notebookCount += 1;
+        // return this.open(model.file);
+        // tslint:disable-next-line: no-any
+        return undefined as any;
+    }
+    protected openedEditor(editor: INotebookEditor): void {
+        this.openedNotebookCount += 1;
+        if (!this.executedEditors.has(editor.file.fsPath)) {
+            editor.executed(this.onExecuted.bind(this));
+        }
+        this.disposables.push(editor.onDidChangeViewState(this.onChangedViewState, this));
+        this.openedEditors.add(editor);
+        editor.closed(this.closedEditor.bind(this));
+        this._onDidOpenNotebookEditor.fire(editor);
+    }
+
+    private closedEditor(editor: INotebookEditor): void {
+        this.openedEditors.delete(editor);
+        this._onDidCloseNotebookEditor.fire(editor);
+    }
+    private onChangedViewState(): void {
+        this._onDidChangeActiveNotebookEditor.fire(this.activeEditor);
+    }
+
+    private onExecuted(editor: INotebookEditor): void {
+        if (editor) {
+            this.executedEditors.add(editor.file.fsPath);
+        }
+    }
+
+    private async onDidOpenNotebookDocument(doc: NotebookDocument | Uri): Promise<void> {
+        const uri = isUri(doc) ? doc : doc.uri;
+        const model = await this.storage.load(uri);
+        // In open method we might be waiting.
+        const editor = this.notebookEditorsByUri.get(uri.toString()) || new NotebookEditor(model);
+        const deferred = this.notebooksWaitingToBeOpenedByUri.get(uri.toString());
+        deferred?.resolve(editor); // NOSONAR
+        if (!isUri(doc)) {
+            // This is where we ensure changes to our models are propagated back to the VSCode model.
+            this.disposables.push(monitorModelCellOutputChangesAndUpdateNotebookDocument(doc, model));
+            this.notebookEditors.set(doc, editor);
+        }
+        this.notebookEditorsByUri.set(uri.toString(), editor);
+    }
+    private async onDidCloseNotebookDocument(doc: NotebookDocument | Uri): Promise<void> {
+        const editor = isUri(doc) ? this.notebookEditorsByUri.get(doc.toString()) : this.notebookEditors.get(doc);
+        if (editor) {
+            editor.dispose();
+            if (editor.model) {
+                editor.model.dispose();
+            }
+        }
+        if (isUri(doc)) {
+            this.notebookEditorsByUri.delete(doc.toString());
+        } else {
+            this.notebookEditors.delete(doc);
+            this.notebookEditorsByUri.delete(doc.uri.toString());
+        }
+    }
+}

src/client/datascience/notebook/serviceRegistry.ts

@@ -4,12 +4,27 @@
 'use strict';
 
 import { IExtensionSingleActivationService } from '../../activation/types';
+import { NativeNotebook } from '../../common/experimentGroups';
+import { IExperimentsManager } from '../../common/types';
 import { IServiceManager } from '../../ioc/types';
+import { INotebookEditorProvider } from '../types';
 import { NotebookContentProvider } from './contentProvider';
 import { NotebookIntegration } from './integration';
+import { NotebookEditorProvider, NotebookEditorProviderActivation } from './notebookEditorProvider';
 import { NotebookKernel } from './notebookKernel';
 
 export function registerTypes(serviceManager: IServiceManager) {
+    // This condition is temporary.
+    // If user belongs to the experiment, then make the necessary changes to package.json.
+    // Once the API is final, we won't need to modify the package.json.
+    if (!serviceManager.get<IExperimentsManager>(IExperimentsManager).inExperiment(NativeNotebook.experiment)) {
+        return;
+    }
+    serviceManager.addSingleton<IExtensionSingleActivationService>(
+        IExtensionSingleActivationService,
+        NotebookEditorProviderActivation
+    );
+    serviceManager.rebindSingleton<INotebookEditorProvider>(INotebookEditorProvider, NotebookEditorProvider);
     serviceManager.addSingleton<NotebookKernel>(NotebookKernel, NotebookKernel);
     serviceManager.addSingleton<IExtensionSingleActivationService>(
         IExtensionSingleActivationService,

src/client/ioc/serviceManager.ts

@@ -91,6 +91,18 @@ export class ServiceManager implements IServiceManager {
         }
     }
 
+    public rebindSingleton<T>(
+        serviceIdentifier: interfaces.ServiceIdentifier<T>,
+        constructor: ClassType<T>,
+        name?: string | number | symbol
+    ): void {
+        if (name) {
+            this.container.rebind<T>(serviceIdentifier).to(constructor).inSingletonScope().whenTargetNamed(name);
+        } else {
+            this.container.rebind<T>(serviceIdentifier).to(constructor).inSingletonScope();
+        }
+    }
+
     public rebindInstance<T>(
         serviceIdentifier: interfaces.ServiceIdentifier<T>,
         instance: T,