Skip to content

Commit e33af82

Browse files
DonJayamanneDonJayamanne
authored
Update VS Code Notebook API (microsoft#13143)
For #12189
1 parent 39f2419 commit e33af82

13 files changed

Lines changed: 291 additions & 154 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -461,7 +461,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
461461
protected async setLaunchingFile(_file: string): Promise<void> {
462462
// For the native editor, use our own file as the path
463463
const notebook = this.getNotebook();
464-
if ((await this.fs.localFileExists(this.file.fsPath)) && notebook) {
464+
if (notebook) {
465465
await notebook.setLaunchingFile(this.file.fsPath);
466466
}
467467
}

‎src/client/datascience/notebook/executionService.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ export class NotebookExecutionService implements INotebookExecutionService {
8989
public async executeAllCells(document: NotebookDocument, token: CancellationToken): Promise<void> {
9090
const stopWatch = new StopWatch();
9191
const notebookAndModel = this.getNotebookAndModel(document);
92-
92+
document.metadata.runState = vscodeNotebookEnums.NotebookRunState.Running;
9393
// Mark all cells as busy (this way there's immediate feedback to users).
9494
// If it does not complete, then restore old state.
9595
const oldCellStates = new WeakMap<NotebookCell, NotebookCellRunState | undefined>();
@@ -113,7 +113,10 @@ export class NotebookExecutionService implements INotebookExecutionService {
113113
}
114114
};
115115
// If we cancel running cells, then restore the state to previous values unless cell has completed.
116-
token.onCancellationRequested(() => document.cells.forEach(restoreOldCellState));
116+
token.onCancellationRequested(() => {
117+
document.metadata.runState = vscodeNotebookEnums.NotebookRunState.Idle;
118+
document.cells.forEach(restoreOldCellState);
119+
});
117120

118121
let executingAPreviousCellHasFailed = false;
119122
await document.cells.reduce((previousPromise, cellToExecute) => {
@@ -137,6 +140,8 @@ export class NotebookExecutionService implements INotebookExecutionService {
137140
return this.executeIndividualCell(notebookAndModel, document, cellToExecute, token, stopWatch);
138141
});
139142
}, Promise.resolve<NotebookCellRunState | undefined>(undefined));
143+
144+
document.metadata.runState = vscodeNotebookEnums.NotebookRunState.Idle;
140145
}
141146
public cancelPendingExecutions(document: NotebookDocument): void {
142147
this.pendingExecutionCancellations.get(document.uri.fsPath)?.forEach((cancellation) => cancellation.cancel()); // NOSONAR

‎src/client/datascience/notebook/notebookEditor.ts‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import {
2222
IStatusProvider
2323
} from '../types';
2424
import { getDefaultCodeLanguage } from './helpers/helpers';
25-
import { INotebookContentProvider, INotebookExecutionService } from './types';
25+
import { INotebookExecutionService } from './types';
2626

2727
export class NotebookEditor implements INotebookEditor {
2828
public readonly type = 'native';
@@ -76,7 +76,6 @@ export class NotebookEditor implements INotebookEditor {
7676
private readonly executionService: INotebookExecutionService,
7777
private readonly commandManager: ICommandManager,
7878
private readonly notebookProvider: INotebookProvider,
79-
private readonly contentProvider: INotebookContentProvider,
8079
private readonly statusProvider: IStatusProvider,
8180
private readonly applicationShell: IApplicationShell,
8281
private readonly configurationService: IConfigurationService,
@@ -153,8 +152,6 @@ export class NotebookEditor implements INotebookEditor {
153152
if (!notebook || this.restartingKernel) {
154153
return;
155154
}
156-
this.restartingKernel = true;
157-
158155
const status = this.statusProvider.set(DataScience.interruptKernelStatus(), true, undefined, undefined);
159156

160157
try {
@@ -178,8 +175,6 @@ export class NotebookEditor implements INotebookEditor {
178175
status.dispose();
179176
traceError(err);
180177
this.applicationShell.showErrorMessage(err);
181-
} finally {
182-
this.restartingKernel = false;
183178
}
184179
}
185180

@@ -230,7 +225,6 @@ export class NotebookEditor implements INotebookEditor {
230225
try {
231226
this.document.metadata.cellRunnable = false;
232227
this.document.metadata.runnable = false;
233-
this.contentProvider.notifyChangesToDocument(this.document);
234228
await notebook.restartKernel(
235229
this.configurationService.getSettings(this.file).datascience.jupyterInterruptTimeout
236230
);
@@ -255,7 +249,6 @@ export class NotebookEditor implements INotebookEditor {
255249
this.restartingKernel = false;
256250
// Restore previous state.
257251
[this.document.metadata.cellRunnable, this.document.metadata.runnable] = [cellRunnable, runnable];
258-
this.contentProvider.notifyChangesToDocument(this.document);
259252
}
260253
}
261254
private async shouldAskForRestart(): Promise<boolean> {

‎src/client/datascience/notebook/notebookEditorProvider.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,6 @@ export class NotebookEditorProvider implements INotebookEditorProvider {
184184
executionService,
185185
this.commandManager,
186186
notebookProvider,
187-
this.contentProvider,
188187
this.statusProvider,
189188
this.appShell,
190189
this.configurationService,

‎src/client/datascience/notebook/notebookKernel.ts‎

Lines changed: 68 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,30 +4,88 @@
44
'use strict';
55

66
import { inject, injectable } from 'inversify';
7-
import { CancellationToken, Uri } from 'vscode';
7+
import { CancellationToken, EventEmitter, Uri } from 'vscode';
88
import type { NotebookCell, NotebookDocument, NotebookKernel as VSCNotebookKernel } from 'vscode-proposed';
9+
import { noop } from '../../common/utils/misc';
910
import { INotebookExecutionService } from './types';
1011

12+
/**
13+
* Cancellation token source that can be cancelled multiple times.
14+
*/
15+
class MultiCancellationTokenSource {
16+
/**
17+
* The cancellation token of this source.
18+
*/
19+
public readonly token: CancellationToken;
20+
private readonly eventEmitter = new EventEmitter<void>();
21+
constructor() {
22+
this.token = {
23+
isCancellationRequested: false,
24+
onCancellationRequested: this.eventEmitter.event.bind(this.eventEmitter)
25+
};
26+
}
27+
public cancel(): void {
28+
this.token.isCancellationRequested = true;
29+
this.eventEmitter.fire();
30+
}
31+
32+
/**
33+
* Dispose object and free resources.
34+
*/
35+
public dispose(): void {
36+
this.eventEmitter.dispose();
37+
}
38+
}
39+
1140
/**
1241
* VSC will use this class to execute cells in a notebook.
1342
* This is where we hookup Jupyter with a Notebook in VSCode.
1443
*/
1544
@injectable()
1645
export class NotebookKernel implements VSCNotebookKernel {
17-
private _preloads: Uri[] = [];
18-
1946
get preloads(): Uri[] {
20-
return this._preloads;
47+
return [];
2148
}
22-
constructor(@inject(INotebookExecutionService) private readonly execution: INotebookExecutionService) {}
2349
public get label(): string {
2450
return 'Jupyter';
2551
}
26-
27-
public async executeCell(document: NotebookDocument, cell: NotebookCell, token: CancellationToken): Promise<void> {
28-
return this.execution.executeCell(document, cell, token);
52+
private cellExecutions = new WeakMap<NotebookCell, MultiCancellationTokenSource>();
53+
private documentExecutions = new WeakMap<NotebookDocument, MultiCancellationTokenSource>();
54+
constructor(@inject(INotebookExecutionService) private readonly execution: INotebookExecutionService) {}
55+
public executeCell(document: NotebookDocument, cell: NotebookCell) {
56+
if (this.cellExecutions.has(cell)) {
57+
return;
58+
}
59+
const source = new MultiCancellationTokenSource();
60+
this.cellExecutions.set(cell, source);
61+
this.execution
62+
.executeCell(document, cell, source.token)
63+
.finally(() => {
64+
if (this.cellExecutions.get(cell) === source) {
65+
this.cellExecutions.delete(cell);
66+
}
67+
})
68+
.catch(noop);
69+
}
70+
public executeAllCells(document: NotebookDocument) {
71+
if (this.documentExecutions.has(document)) {
72+
return;
73+
}
74+
const source = new MultiCancellationTokenSource();
75+
this.documentExecutions.set(document, source);
76+
this.execution
77+
.executeAllCells(document, source.token)
78+
.finally(() => {
79+
if (this.documentExecutions.get(document) === source) {
80+
this.documentExecutions.delete(document);
81+
}
82+
})
83+
.catch(noop);
84+
}
85+
public cancelCellExecution(_document: NotebookDocument, cell: NotebookCell) {
86+
this.cellExecutions.get(cell)?.cancel(); // NOSONAR
2987
}
30-
public async executeAllCells(document: NotebookDocument, token: CancellationToken): Promise<void> {
31-
return this.execution.executeAllCells(document, token);
88+
public cancelAllCellsExecution(document: NotebookDocument) {
89+
this.documentExecutions.get(document)?.cancel(); // NOSONAR
3290
}
3391
}

‎src/test/datascience/notebook/cellOutput.ds.test.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,6 @@ suite('DataScience - VSCode Notebook - (fake execution) (Clearing Output)', func
142142

143143
// Once we execute the cell, the execution count & output should be cleared.
144144
await commands.executeCommand('notebook.cell.execute');
145-
146145
await waitForExecutionOrderInVSCCell(vscCell, undefined);
147146
await waitForVSCCellHasEmptyOutput(vscCell);
148147
await waitForVSCCellIsRunning(vscCell);

‎src/test/datascience/notebook/executionErrors.ds.test.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ import {
1515
INotebookEditorProvider,
1616
INotebookProvider
1717
} from '../../../client/datascience/types';
18-
import { IExtensionTestApi } from '../../common';
18+
import { IExtensionTestApi, waitForCondition } from '../../common';
1919
import { initialize, initializeTest } from '../../initialize';
2020
import { canRunTests, closeNotebooksAndCleanUpAfterTests, insertPythonCellAndWait } from './helper';
2121

@@ -61,7 +61,7 @@ suite('DataScience - VSCode Notebook - Errors in Execution', function () {
6161
when(notebook.executeObservable(anything(), anything(), anything(), anything(), anything())).thenThrow(error);
6262
await commands.executeCommand('notebook.execute');
6363

64-
assert.isTrue(handleErrorStub.calledOnce);
64+
await waitForCondition(async () => handleErrorStub.calledOnce, 5_000, 'handleError not called');
6565
assert.isTrue(handleErrorStub.calledOnceWithExactly(error));
6666
});
6767
test('Errors thrown in cell execution (jupyter results) are handled by error handler', async () => {
@@ -80,7 +80,7 @@ suite('DataScience - VSCode Notebook - Errors in Execution', function () {
8080
// Execute cells (it should throw an error).
8181
await commands.executeCommand('notebook.execute');
8282

83-
assert.isTrue(handleErrorStub.calledOnce);
83+
await waitForCondition(async () => handleErrorStub.calledOnce, 5_000, 'handleError not called');
8484
assert.isTrue(handleErrorStub.calledOnceWithExactly(error));
8585
});
8686
});

‎src/test/datascience/notebook/helper.ts‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -298,10 +298,16 @@ export function assertHasTextOutputInVSCode(cell: NotebookCell, text: string, in
298298
}
299299
return true;
300300
}
301-
export async function waitForTextOutputInVSCode(cell: NotebookCell, text: string, index: number, isExactMatch = true) {
301+
export async function waitForTextOutputInVSCode(
302+
cell: NotebookCell,
303+
text: string,
304+
index: number,
305+
isExactMatch = true,
306+
timeout = 1_000
307+
) {
302308
await waitForCondition(
303309
async () => assertHasTextOutputInVSCode(cell, text, index, isExactMatch),
304-
1_000,
310+
timeout,
305311
`Output does not contain provided text '${text}' for Cell ${cell.notebook.cells.indexOf(cell) + 1}`
306312
);
307313
}

0 commit comments

Comments
 (0)