Handling Save as of notebooks when using new VSC Notebook API (#11876)

For #10496
* Partial implementation of `Save As`.
* There's more to this, but wanted to get in what I had.

Final implementation details outlined here:
* Based on discussions with team, https://github.com/microsoft/vscode-python/issues/11877
* We will need to wait for some VSCode API to be made available before we can complete this.

Author: DonJayamanne

src/client/datascience/interactive-common/interactiveWindowTypes.ts

@@ -345,6 +345,10 @@ export interface INotebookModelChange {
 export interface INotebookModelSaved extends INotebookModelChange {
     kind: 'save';
 }
+export interface INotebookModelSavedAs extends INotebookModelChange {
+    kind: 'saveAs';
+    target: Uri;
+}
 
 export interface INotebookModelRemoveAllChange extends INotebookModelChange {
     kind: 'remove_all';
@@ -461,6 +465,7 @@ export interface INotebookModelVersionChange extends INotebookModelChange {
 
 export type NotebookModelChange =
     | INotebookModelSaved
+    | INotebookModelSavedAs
     | INotebookModelModifyChange
     | INotebookModelRemoveAllChange
     | INotebookModelClearChange

src/client/datascience/interactive-common/notebookProvider.ts

@@ -10,6 +10,7 @@ import { IFileSystem } from '../../common/platform/types';
 import { IDisposableRegistry, Resource } from '../../common/types';
 import { noop } from '../../common/utils/misc';
 import { Identifiers } from '../constants';
+import { INotebookStorageProvider } from '../interactive-ipynb/notebookStorageProvider';
 import {
     ConnectNotebookProviderOptions,
     GetNotebookOptions,
@@ -37,9 +38,11 @@ export class NotebookProvider implements INotebookProvider {
         @inject(IDisposableRegistry) disposables: IDisposableRegistry,
         @inject(IRawNotebookProvider) private readonly rawNotebookProvider: IRawNotebookProvider,
         @inject(IJupyterNotebookProvider) private readonly jupyterNotebookProvider: IJupyterNotebookProvider,
-        @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService
+        @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService,
+        @inject(INotebookStorageProvider) storageProvider: INotebookStorageProvider
     ) {
         disposables.push(editorProvider.onDidCloseNotebookEditor(this.onDidCloseNotebookEditor, this));
+        disposables.push(storageProvider.onSavedAs(this.onSavedAs, this));
         disposables.push(
             interactiveWindowProvider.onDidChangeActiveInteractiveWindow(this.checkAndDisposeNotebook, this)
         );
@@ -156,6 +159,15 @@ export class NotebookProvider implements INotebookProvider {
         }
     }
 
+    private async onSavedAs(e: { new: Uri; old: Uri }) {
+        // Swap the Uris when a notebook is saved as a different file.
+        const notebookPromise = this.notebooks.get(e.old.toString());
+        if (notebookPromise) {
+            this.notebooks.set(e.new.toString(), notebookPromise);
+            this.notebooks.delete(e.old.toString());
+        }
+    }
+
     /**
      * Interactive windows have just one window.
      * When that it closed, just close all of the notebooks associated with interactive windows.

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

@@ -29,6 +29,7 @@ import { captureTelemetry, sendTelemetryEvent } from '../../telemetry';
 import { Telemetry } from '../constants';
 import { NotebookModelChange } from '../interactive-common/interactiveWindowTypes';
 import { INotebookEditor, INotebookEditorProvider, INotebookModel } from '../types';
+import { isUntitled } from './nativeEditorStorage';
 import { INotebookStorageProvider } from './notebookStorageProvider';
 
 // Class that is registered as the custom editor provider for notebooks. VS code will call into this class when
@@ -265,10 +266,8 @@ export class NativeEditorProvider
     }
 
     private async getNextNewNotebookUri(): Promise<Uri> {
-        // tslint:disable-next-line: no-suspicious-comment
-        // TODO: This will not work, if we close an untitled document.
         // See if we have any untitled storage already
-        const untitledStorage = Array.from(this.models.values()).filter((model) => model?.file?.scheme === 'untitled');
+        const untitledStorage = Array.from(this.models.values()).filter((model) => model && isUntitled(model));
         // Just use the length (don't bother trying to fill in holes). We never remove storage objects from
         // our map, so we'll keep creating new untitled notebooks.
         const fileName = `${localize.DataScience.untitledNotebookFileName()}-${untitledStorage.length + 1}.ipynb`;

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 { isFileNotFoundError } from '../../common/platform/errors';
 import { sendTelemetryEvent } from '../../telemetry';
 import { pruneCell } from '../common';
 
@@ -34,6 +35,13 @@ interface INativeEditorStorageState {
     notebookJson: Partial<nbformat.INotebookContent>;
 }
 
+function isUntitledFile(file?: Uri) {
+    return file?.scheme === 'untitled';
+}
+export function isUntitled(model?: INotebookModel): boolean {
+    return isUntitledFile(model?.file);
+}
+
 class NativeEditorNotebookModel implements INotebookModel {
     public get onDidDispose() {
         return this._disposed.event;
@@ -52,14 +60,17 @@ class NativeEditorNotebookModel implements INotebookModel {
     }
 
     public get isUntitled(): boolean {
-        return this.file.scheme === 'untitled';
+        return isUntitled(this);
     }
     public get cells(): ICell[] {
         return this._state.cells;
     }
     public get onDidEdit(): Event<NotebookModelChange> {
         return this._editEventEmitter.event;
     }
+    public get metadata(): nbformat.INotebookMetadata | undefined {
+        return this._state.notebookJson.metadata;
+    }
     private _disposed = new EventEmitter<void>();
     private _isDisposed?: boolean;
     private _changedEmitter = new EventEmitter<NotebookModelChange>();
@@ -96,7 +107,6 @@ class NativeEditorNotebookModel implements INotebookModel {
         this._isDisposed = true;
         this._disposed.fire();
     }
-
     public clone(file: Uri) {
         return new NativeEditorNotebookModel(
             file,
@@ -115,9 +125,6 @@ class NativeEditorNotebookModel implements INotebookModel {
     public async undoEdits(edits: readonly NotebookModelChange[]): Promise<void> {
         edits.forEach((e) => this.update({ ...e, source: 'undo' }));
     }
-    public get metadata(): nbformat.INotebookMetadata | undefined {
-        return this._state.notebookJson.metadata;
-    }
 
     public getContent(): string {
         return this.generateNotebookContent();
@@ -182,13 +189,18 @@ class NativeEditorNotebookModel implements INotebookModel {
             case 'save':
                 this._state.saveChangeCount = this._state.changeCount;
                 break;
+            case 'saveAs':
+                this._state.saveChangeCount = this._state.changeCount;
+                this._state.changeCount = this._state.saveChangeCount = 0;
+                this._state.file = change.target;
+                break;
             default:
                 break;
         }
 
         // Dirty state comes from undo. At least VS code will track it that way. However
         // skip file changes as we don't forward those to VS code
-        if (change.kind !== 'save') {
+        if (change.kind !== 'save' && change.kind !== 'saveAs') {
             this._state.changeCount += 1;
         }
 
@@ -433,6 +445,10 @@ class NativeEditorNotebookModel implements INotebookModel {
 
 @injectable()
 export class NativeEditorStorage implements INotebookStorage {
+    public get onSavedAs(): Event<{ new: Uri; old: Uri }> {
+        return this.savedAs.event;
+    }
+    private readonly savedAs = new EventEmitter<{ new: Uri; old: Uri }>();
     constructor(
         @inject(IJupyterExecution) private jupyterExecution: IJupyterExecution,
         @inject(IFileSystem) private fileSystem: IFileSystem,
@@ -442,25 +458,35 @@ export class NativeEditorStorage implements INotebookStorage {
         @inject(IMemento) @named(WORKSPACE_MEMENTO) private localStorage: Memento
     ) {}
     private static isUntitledFile(file: Uri) {
-        return file.scheme === 'untitled';
+        return isUntitledFile(file);
     }
 
     public async load(file: Uri, possibleContents?: string): Promise<INotebookModel> {
         return this.loadFromFile(file, possibleContents);
     }
     public async save(model: INotebookModel, _cancellation: CancellationToken): Promise<void> {
-        await this.saveAs(model, model.file);
+        const contents = model.getContent();
+        await this.fileSystem.writeFile(model.file.fsPath, contents, 'utf-8');
+        model.update({
+            source: 'user',
+            kind: 'save',
+            oldDirty: model.isDirty,
+            newDirty: false
+        });
     }
 
     public async saveAs(model: INotebookModel, file: Uri): Promise<void> {
+        const old = model.file;
         const contents = model.getContent();
         await this.fileSystem.writeFile(file.fsPath, contents, 'utf-8');
         model.update({
             source: 'user',
-            kind: 'save',
+            kind: 'saveAs',
             oldDirty: model.isDirty,
-            newDirty: false
+            newDirty: false,
+            target: file
         });
+        this.savedAs.fire({ new: file, old });
     }
     public async backup(model: INotebookModel, cancellation: CancellationToken): Promise<void> {
         // Should send to extension context storage path
@@ -657,7 +683,10 @@ export class NativeEditorStorage implements INotebookStorage {
                 return data.contents;
             }
         } catch (exc) {
-            traceError(`Exception reading from temporary storage for ${key}`, exc);
+            // No need to log error if file doesn't exist.
+            if (!isFileNotFoundError(exc)) {
+                traceError(`Exception reading from temporary storage for ${key}`, exc);
+            }
         }
     }
 

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

@@ -4,20 +4,29 @@
 'use strict';
 
 import { inject, injectable } from 'inversify';
-import { CancellationToken, CancellationTokenSource, Uri, WorkspaceConfiguration } from 'vscode';
+import { CancellationToken, CancellationTokenSource, EventEmitter, Uri, WorkspaceConfiguration } from 'vscode';
 import { IWorkspaceService } from '../../common/application/types';
 import { IDisposable, IDisposableRegistry } from '../../common/types';
+import { DataScience } from '../../common/utils/localize';
 import { noop } from '../../common/utils/misc';
 import { INotebookModel, INotebookStorage } from '../types';
+import { isUntitled } from './nativeEditorStorage';
 
 // tslint:disable-next-line:no-require-imports no-var-requires
 const debounce = require('lodash/debounce') as typeof import('lodash/debounce');
 
 export const INotebookStorageProvider = Symbol.for('INotebookStorageProvider');
-export interface INotebookStorageProvider extends INotebookStorage {}
+export interface INotebookStorageProvider extends INotebookStorage {
+    createNew(contents?: string): Promise<INotebookModel>;
+}
 @injectable()
 export class NotebookStorageProvider implements INotebookStorageProvider {
+    public get onSavedAs() {
+        return this._savedAs.event;
+    }
+    private readonly _savedAs = new EventEmitter<{ new: Uri; old: Uri }>();
     private readonly storageAndModels = new Map<string, Promise<INotebookModel>>();
+    private models = new Set<INotebookModel>();
     private readonly disposables: IDisposable[] = [];
     private readonly _autoSaveNotebookInHotExitFile = new WeakMap<INotebookModel, Function>();
     constructor(
@@ -26,13 +35,16 @@ export class NotebookStorageProvider implements INotebookStorageProvider {
         @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService
     ) {
         disposables.push(this);
+        disposables.push(storage.onSavedAs((e) => this._savedAs.fire(e)));
     }
     public save(model: INotebookModel, cancellation: CancellationToken) {
         return this.storage.save(model, cancellation);
     }
     public async saveAs(model: INotebookModel, targetResource: Uri) {
+        const oldUri = model.file;
         await this.storage.saveAs(model, targetResource);
         this.trackModel(model);
+        this.storageAndModels.delete(oldUri.toString());
         this.storageAndModels.set(targetResource.toString(), Promise.resolve(model));
     }
     public backup(model: INotebookModel, cancellation: CancellationToken) {
@@ -52,11 +64,33 @@ export class NotebookStorageProvider implements INotebookStorageProvider {
             this.disposables.shift()?.dispose(); // NOSONAR
         }
     }
+
+    public async createNew(contents?: string): Promise<INotebookModel> {
+        // Create a new URI for the dummy file using our root workspace path
+        const uri = await this.getNextNewNotebookUri();
+        return this.load(uri, contents);
+    }
+
+    private async getNextNewNotebookUri(): Promise<Uri> {
+        // tslint:disable-next-line: no-suspicious-comment
+        // TODO: This will not work, if we close an untitled document.
+        // See if we have any untitled storage already
+        const untitledStorage = Array.from(this.models.values()).filter(isUntitled);
+        // Just use the length (don't bother trying to fill in holes). We never remove storage objects from
+        // our map, so we'll keep creating new untitled notebooks.
+        const fileName = `${DataScience.untitledNotebookFileName()}-${untitledStorage.length + 1}.ipynb`;
+        const fileUri = Uri.file(fileName);
+        // Turn this back into an untitled
+        return fileUri.with({ scheme: 'untitled', path: fileName });
+    }
+
     private trackModel(model: INotebookModel) {
         this.disposables.push(model);
+        this.models.add(model);
         // When a model is no longer used, ensure we remove it from the cache.
         model.onDidDispose(
             () => {
+                this.models.delete(model);
                 this.storageAndModels.delete(model.file.toString());
                 this._autoSaveNotebookInHotExitFile.delete(model);
             },

src/client/datascience/notebook/executionHelpers.ts

@@ -89,7 +89,25 @@ export function monitorModelCellOutputChangesAndUpdateNotebookDocument(
     document: NotebookDocument,
     model: INotebookModel
 ): IDisposable {
-    return model.changed((change) => {
+    let wasUntitledNotebook = model.isUntitled;
+    let stopSyncingOutput = false;
+    const disposable = model.changed((change) => {
+        if (stopSyncingOutput) {
+            return;
+        }
+        if (change.kind === 'saveAs') {
+            if (wasUntitledNotebook) {
+                wasUntitledNotebook = false;
+                // User saved untitled file as a real file.
+                return;
+            } else {
+                // Ok, user save a normal notebook as another name.
+                // Stop monitoring changes.
+                stopSyncingOutput = true;
+                disposable.dispose();
+                return;
+            }
+        }
         // We're only interested in updates to cells.
         if (change.kind !== 'modify') {
             return;
@@ -117,6 +135,8 @@ export function monitorModelCellOutputChangesAndUpdateNotebookDocument(
             uiCellToUpdate.outputs = newOutput;
         }
     });
+
+    return disposable;
 }
 
 /**

src/client/datascience/notebook/notebookEditorProvider.ts

@@ -81,6 +81,17 @@ export class NotebookEditorProvider implements INotebookEditorProvider {
         this.disposables.push(notebook.onDidOpenNotebookDocument(this.onDidOpenNotebookDocument, this));
         this.disposables.push(notebook.onDidCloseNotebookDocument(this.onDidCloseNotebookDocument, this));
 
+        // Swap the uris.
+        this.disposables.push(
+            this.storage.onSavedAs((e) => {
+                const savedEditor = this.notebookEditorsByUri.get(e.old.toString());
+                if (savedEditor) {
+                    this.notebookEditorsByUri.delete(e.old.toString());
+                    this.notebookEditorsByUri.set(e.new.toString(), savedEditor);
+                }
+            })
+        );
+
         // 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');
@@ -114,18 +125,15 @@ export class NotebookEditorProvider implements INotebookEditorProvider {
         // 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);
+    public async createNew(contents?: string): Promise<INotebookEditor> {
+        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;
+
+        return this.open(model.file);
     }
     protected openedEditor(editor: INotebookEditor): void {
         this.openedNotebookCount += 1;