Skip to content

Commit 06af191

Browse files
rchiodorchiodo
authored
Support hot exit storage for untitled notebooks (#11891)
* Basic tracking * Fixup untitled number increment * Add unit test * Add more tests * Handle backup for untitled and fixup tests * Rework tests to verify storage and accomodate weird ness in code * Dont open untitled files if they weren't changed * Add news entry * Fix sonar problem * Add comment
1 parent 371f036 commit 06af191

14 files changed

+673
-230
lines changed

‎news/2 Fixes/11711.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Reopen all notebooks when rerunning the extension (including untitled ones).

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import {
2828
IJupyterVariable,
2929
IJupyterVariablesRequest,
3030
IJupyterVariablesResponse,
31+
INotebookModel,
3132
KernelSocketOptions
3233
} from '../types';
3334
import { BaseReduxActionPayload } from './types';
@@ -340,6 +341,7 @@ export interface INotebookModelChange {
340341
oldDirty: boolean;
341342
newDirty: boolean;
342343
source: 'undo' | 'user' | 'redo';
344+
model?: INotebookModel;
343345
}
344346

345347
export interface INotebookModelSaved extends INotebookModelChange {

‎src/client/datascience/interactive-ipynb/nativeEditorOldWebView.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,7 @@ export class NativeEditorOldWebView extends NativeEditor {
200200
// Skip doing this if auto save is enabled.
201201
const filesConfig = this.workspaceService.getConfiguration('files', this.file);
202202
const autoSave = filesConfig.get('autoSave', 'off');
203-
if (autoSave === 'off') {
203+
if (autoSave === 'off' || this.isUntitled) {
204204
await this.storage.backup(this.model, new CancellationTokenSource().token);
205205
}
206206
this.commandManager.executeCommand(Commands.OpenNotebookNonCustomEditor, this.model.file).then(noop, noop);

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

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ import { captureTelemetry, sendTelemetryEvent } from '../../telemetry';
2929
import { Telemetry } from '../constants';
3030
import { NotebookModelChange } from '../interactive-common/interactiveWindowTypes';
3131
import { INotebookEditor, INotebookEditorProvider, INotebookModel } from '../types';
32-
import { isUntitled } from './nativeEditorStorage';
32+
import { getNextUntitledCounter } from './nativeEditorStorage';
3333
import { INotebookStorageProvider } from './notebookStorageProvider';
3434

3535
// Class that is registered as the custom editor provider for notebooks. VS code will call into this class when
@@ -76,6 +76,7 @@ export class NativeEditorProvider
7676
private notebookCount: number = 0;
7777
private openedNotebookCount: number = 0;
7878
private _id = uuid();
79+
private untitledCounter = 1;
7980
constructor(
8081
@inject(IServiceContainer) protected readonly serviceContainer: IServiceContainer,
8182
@inject(IAsyncDisposableRegistry) protected readonly asyncRegistry: IAsyncDisposableRegistry,
@@ -191,14 +192,17 @@ export class NativeEditorProvider
191192
// Update number of notebooks in the workspace
192193
this.notebookCount += 1;
193194

194-
// Set these contents into the storage before the file opens
195-
await this.loadModel(uri, contents);
195+
// Set these contents into the storage before the file opens. Make sure not
196+
// load from the memento storage though as this is an entirely brand new file.
197+
await this.loadModel(uri, contents, true);
196198

197199
return this.open(uri);
198200
}
199201

200-
public loadModel(file: Uri, contents?: string) {
201-
return this.storage.load(file, contents).then((m) => {
202+
public loadModel(file: Uri, contents?: string, skipDirtyContents?: boolean) {
203+
// Every time we load a new untitled file, up the counter past the max value for this counter
204+
this.untitledCounter = getNextUntitledCounter(file, this.untitledCounter);
205+
return this.storage.load(file, contents, skipDirtyContents).then((m) => {
202206
this.trackModel(m);
203207
return m;
204208
});
@@ -266,11 +270,8 @@ export class NativeEditorProvider
266270
}
267271

268272
private async getNextNewNotebookUri(): Promise<Uri> {
269-
// See if we have any untitled storage already
270-
const untitledStorage = Array.from(this.models.values()).filter((model) => model && isUntitled(model));
271-
// Just use the length (don't bother trying to fill in holes). We never remove storage objects from
272-
// our map, so we'll keep creating new untitled notebooks.
273-
const fileName = `${localize.DataScience.untitledNotebookFileName()}-${untitledStorage.length + 1}.ipynb`;
273+
// Just use the current counter. Counter will be incremented after actually opening a file.
274+
const fileName = `${localize.DataScience.untitledNotebookFileName()}-${this.untitledCounter}.ipynb`;
274275
const fileUri = Uri.file(fileName);
275276
// Turn this back into an untitled
276277
return fileUri.with({ scheme: 'untitled', path: fileName });

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

Lines changed: 58 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,27 @@ export function isUntitled(model?: INotebookModel): boolean {
4242
return isUntitledFile(model?.file);
4343
}
4444

45-
class NativeEditorNotebookModel implements INotebookModel {
45+
export function getNextUntitledCounter(file: Uri | undefined, currentValue: number): number {
46+
if (file && isUntitledFile(file)) {
47+
const basename = path.basename(file.fsPath, 'ipynb');
48+
const extname = path.extname(file.fsPath);
49+
if (extname.toLowerCase() === '.ipynb') {
50+
// See if ends with -<n>
51+
const match = /.*-(\d+)/.exec(basename);
52+
if (match && match[1]) {
53+
const fileValue = parseInt(match[1], 10);
54+
if (fileValue) {
55+
return Math.max(currentValue, fileValue + 1);
56+
}
57+
}
58+
}
59+
}
60+
61+
return currentValue;
62+
}
63+
64+
// Exported for test mocks
65+
export class NativeEditorNotebookModel implements INotebookModel {
4666
public get onDidDispose() {
4767
return this._disposed.event;
4868
}
@@ -71,6 +91,9 @@ class NativeEditorNotebookModel implements INotebookModel {
7191
public get metadata(): nbformat.INotebookMetadata | undefined {
7292
return this._state.notebookJson.metadata;
7393
}
94+
public get id() {
95+
return this._id;
96+
}
7497
private _disposed = new EventEmitter<void>();
7598
private _isDisposed?: boolean;
7699
private _changedEmitter = new EventEmitter<NotebookModelChange>();
@@ -83,6 +106,8 @@ class NativeEditorNotebookModel implements INotebookModel {
83106
notebookJson: {}
84107
};
85108

109+
private _id = uuid();
110+
86111
constructor(
87112
file: Uri,
88113
cells: ICell[],
@@ -148,7 +173,7 @@ class NativeEditorNotebookModel implements INotebookModel {
148173

149174
// Forward onto our listeners if necessary
150175
if (changed || this.isDirty !== oldDirty) {
151-
this._changedEmitter.fire({ ...change, newDirty: this.isDirty, oldDirty });
176+
this._changedEmitter.fire({ ...change, newDirty: this.isDirty, oldDirty, model: this });
152177
}
153178
// Slightly different for the event we send to VS code. Skip version and file changes. Only send user events.
154179
if ((changed || this.isDirty !== oldDirty) && change.kind !== 'version' && change.source === 'user') {
@@ -461,8 +486,8 @@ export class NativeEditorStorage implements INotebookStorage {
461486
return isUntitledFile(file);
462487
}
463488

464-
public async load(file: Uri, possibleContents?: string): Promise<INotebookModel> {
465-
return this.loadFromFile(file, possibleContents);
489+
public async load(file: Uri, possibleContents?: string, skipDirtyContents?: boolean): Promise<INotebookModel> {
490+
return this.loadFromFile(file, possibleContents, skipDirtyContents);
466491
}
467492
public async save(model: INotebookModel, _cancellation: CancellationToken): Promise<void> {
468493
const contents = model.getContent();
@@ -473,6 +498,7 @@ export class NativeEditorStorage implements INotebookStorage {
473498
oldDirty: model.isDirty,
474499
newDirty: false
475500
});
501+
this.clearHotExit(model.file).ignoreErrors();
476502
}
477503

478504
public async saveAs(model: INotebookModel, file: Uri): Promise<void> {
@@ -487,6 +513,7 @@ export class NativeEditorStorage implements INotebookStorage {
487513
target: file
488514
});
489515
this.savedAs.fire({ new: file, old });
516+
this.clearHotExit(old).ignoreErrors();
490517
}
491518
public async backup(model: INotebookModel, cancellation: CancellationToken): Promise<void> {
492519
// Should send to extension context storage path
@@ -508,16 +535,23 @@ export class NativeEditorStorage implements INotebookStorage {
508535

509536
return this.writeToStorage(filePath, specialContents, cancelToken);
510537
}
538+
539+
private async clearHotExit(file: Uri): Promise<void> {
540+
const key = this.getStorageKey(file);
541+
const filePath = this.getHashedFileName(key);
542+
return this.writeToStorage(filePath, undefined);
543+
}
544+
511545
private async writeToStorage(filePath: string, contents?: string, cancelToken?: CancellationToken): Promise<void> {
512546
try {
513547
if (!cancelToken?.isCancellationRequested) {
514548
if (contents) {
515549
await this.fileSystem.createDirectory(path.dirname(filePath));
516550
if (!cancelToken?.isCancellationRequested) {
517-
return this.fileSystem.writeFile(filePath, contents);
551+
return await this.fileSystem.writeFile(filePath, contents);
518552
}
519-
} else {
520-
return this.fileSystem.deleteFile(filePath);
553+
} else if (await this.fileSystem.fileExists(filePath)) {
554+
return await this.fileSystem.deleteFile(filePath);
521555
}
522556
}
523557
} catch (exc) {
@@ -561,15 +595,24 @@ export class NativeEditorStorage implements INotebookStorage {
561595
noop();
562596
}
563597
}
564-
private async loadFromFile(file: Uri, possibleContents?: string): Promise<INotebookModel> {
598+
private async loadFromFile(
599+
file: Uri,
600+
possibleContents?: string,
601+
skipDirtyContents?: boolean
602+
): Promise<INotebookModel> {
565603
try {
566604
// Attempt to read the contents if a viable file
567605
const contents = NativeEditorStorage.isUntitledFile(file)
568606
? possibleContents
569607
: await this.fileSystem.readFile(file.fsPath);
570608

609+
// If skipping dirty contents, delete the dirty hot exit file now
610+
if (skipDirtyContents) {
611+
await this.clearHotExit(file);
612+
}
613+
571614
// See if this file was stored in storage prior to shutdown
572-
const dirtyContents = await this.getStoredContents(file);
615+
const dirtyContents = skipDirtyContents ? undefined : await this.getStoredContents(file);
573616
if (dirtyContents) {
574617
// This means we're dirty. Indicate dirty and load from this content
575618
return this.loadContents(file, dirtyContents, true);
@@ -668,18 +711,13 @@ export class NativeEditorStorage implements INotebookStorage {
668711
const contents = await this.fileSystem.readFile(filePath);
669712
const data = JSON.parse(contents);
670713
// Check whether the file has been modified since the last time the contents were saved.
671-
if (
672-
data &&
673-
data.lastModifiedTimeMs &&
674-
!NativeEditorStorage.isUntitledFile(file) &&
675-
file.scheme === 'file'
676-
) {
714+
if (data && data.lastModifiedTimeMs && file.scheme === 'file') {
677715
const stat = await this.fileSystem.stat(file.fsPath);
678716
if (stat.mtime > data.lastModifiedTimeMs) {
679717
return;
680718
}
681719
}
682-
if (data && !NativeEditorStorage.isUntitledFile(file) && data.contents) {
720+
if (data && data.contents) {
683721
return data.contents;
684722
}
685723
} catch (exc) {
@@ -700,28 +738,23 @@ export class NativeEditorStorage implements INotebookStorage {
700738
}
701739

702740
// Check whether the file has been modified since the last time the contents were saved.
703-
if (
704-
data &&
705-
data.lastModifiedTimeMs &&
706-
!NativeEditorStorage.isUntitledFile(file) &&
707-
file.scheme === 'file'
708-
) {
741+
if (data && data.lastModifiedTimeMs && file.scheme === 'file') {
709742
const stat = await this.fileSystem.stat(file.fsPath);
710743
if (stat.mtime > data.lastModifiedTimeMs) {
711744
return;
712745
}
713746
}
714-
if (data && !NativeEditorStorage.isUntitledFile(file) && data.contents) {
747+
if (data && data.contents) {
715748
return data.contents;
716749
}
717750
} catch {
718751
noop();
719752
}
720753
}
721754

722-
private async getStoredContentsFromLocalStorage(file: Uri, key: string): Promise<string | undefined> {
755+
private async getStoredContentsFromLocalStorage(_file: Uri, key: string): Promise<string | undefined> {
723756
const workspaceData = this.localStorage.get<string>(key);
724-
if (workspaceData && !NativeEditorStorage.isUntitledFile(file)) {
757+
if (workspaceData) {
725758
// Make sure to clear so we don't use this again.
726759
this.localStorage.update(key, undefined);
727760

‎src/client/datascience/interactive-ipynb/nativeEditorViewTracker.ts‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
import { inject, injectable, named } from 'inversify';
2+
import { Memento, Uri } from 'vscode';
3+
import { IExtensionSingleActivationService } from '../../activation/types';
4+
import { IDisposableRegistry, IMemento, WORKSPACE_MEMENTO } from '../../common/types';
5+
import { INotebookEditor, INotebookEditorProvider } from '../types';
6+
import { isUntitled } from './nativeEditorStorage';
7+
8+
const MEMENTO_KEY = 'nativeEditorViewTracking';
9+
/**
10+
* This class tracks opened notebooks and stores the list of files in a memento. On next activation
11+
* this list of files is then opened.
12+
* Untitled files are tracked too, but they should only open if they're dirty.
13+
*/
14+
@injectable()
15+
export class NativeEditorViewTracker implements IExtensionSingleActivationService {
16+
constructor(
17+
@inject(INotebookEditorProvider) private readonly editorProvider: INotebookEditorProvider,
18+
@inject(IMemento) @named(WORKSPACE_MEMENTO) private readonly workspaceMemento: Memento,
19+
@inject(IDisposableRegistry) disposableRegistry: IDisposableRegistry
20+
) {
21+
disposableRegistry.push(editorProvider.onDidOpenNotebookEditor(this.onOpenedEditor.bind(this)));
22+
disposableRegistry.push(editorProvider.onDidCloseNotebookEditor(this.onClosedEditor.bind(this)));
23+
}
24+
25+
public async activate(): Promise<void> {
26+
// On activate get the list and eliminate any dupes that might have snuck in.
27+
const set = new Set<string>(this.workspaceMemento.get<string[]>(MEMENTO_KEY) || []);
28+
await this.workspaceMemento.update(MEMENTO_KEY, undefined);
29+
30+
// Then open each one.
31+
set.forEach((l) => {
32+
const uri = Uri.parse(l);
33+
if (uri) {
34+
this.editorProvider.open(uri).ignoreErrors();
35+
}
36+
});
37+
}
38+
39+
private onOpenedEditor(editor: INotebookEditor) {
40+
// Save this as a file that should be reopened in this workspace
41+
const list = this.workspaceMemento.get<string[]>(MEMENTO_KEY) || [];
42+
const fileKey = editor.file.toString();
43+
44+
// Skip untitled files. They have to be changed first.
45+
if (!list.includes(fileKey) && (!isUntitled(editor.model) || editor.model?.isDirty)) {
46+
this.workspaceMemento.update(MEMENTO_KEY, [...list, fileKey]);
47+
} else if (isUntitled(editor.model) && editor.model) {
48+
editor.model.changed(this.onUntitledChanged.bind(this, editor.file));
49+
}
50+
}
51+
52+
private onUntitledChanged(file: Uri) {
53+
const list = this.workspaceMemento.get<string[]>(MEMENTO_KEY) || [];
54+
const fileKey = file.toString();
55+
if (!list.includes(fileKey)) {
56+
this.workspaceMemento.update(MEMENTO_KEY, [...list, fileKey]);
57+
}
58+
}
59+
60+
private onClosedEditor(editor: INotebookEditor) {
61+
// Save this as a file that should not be reopened in this workspace if this is the
62+
// last editor for this file
63+
const fileKey = editor.file.toString();
64+
if (!this.editorProvider.editors.find((e) => e.file.toString() === fileKey && e !== editor)) {
65+
const list = this.workspaceMemento.get<string[]>(MEMENTO_KEY) || [];
66+
this.workspaceMemento.update(
67+
MEMENTO_KEY,
68+
list.filter((e) => e !== fileKey)
69+
);
70+
}
71+
}
72+
}

0 commit comments

Comments
 (0)