Skip to content

Commit b4eed44

Browse files
DonJayamanneDonJayamanne
committed
Use WorkspaceEdit instead of NotebookEditor
1 parent 9a84217 commit b4eed44

File tree

4 files changed

+133
-106
lines changed

4 files changed

+133
-106
lines changed

‎src/client/datascience/jupyter/kernels/cellExecution.ts

Lines changed: 35 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,14 @@
55

66
import { nbformat } from '@jupyterlab/coreutils';
77
import type { KernelMessage } from '@jupyterlab/services/lib/kernel/messages';
8-
import { CancellationToken, CellOutputKind, CellStreamOutput, NotebookCell, NotebookCellRunState } from 'vscode';
8+
import {
9+
CancellationToken,
10+
CellOutputKind,
11+
CellStreamOutput,
12+
NotebookCell,
13+
NotebookCellRunState,
14+
WorkspaceEdit
15+
} from 'vscode';
916
import type { NotebookEditor as VSCNotebookEditor } from '../../../../../types/vscode-proposed';
1017
import { concatMultilineString, formatStreamText } from '../../../../datascience-ui/common';
1118
import { IApplicationShell, IVSCodeNotebook } from '../../../common/application/types';
@@ -164,24 +171,32 @@ export class CellExecution {
164171

165172
private completedSuccessfully() {
166173
this.sendPerceivedCellExecute();
174+
let statusMessage = '';
167175
// If we requested a cancellation, then assume it did not even run.
168176
// If it did, then we'd get an interrupt error in the output.
169-
this.cell.metadata.runState = this.token.isCancellationRequested
177+
let runState = this.token.isCancellationRequested
170178
? vscodeNotebookEnums.NotebookCellRunState.Idle
171179
: vscodeNotebookEnums.NotebookCellRunState.Success;
172180

173-
this.cell.metadata.statusMessage = '';
174-
this.cell.metadata.lastRunDuration = this.stopWatch.elapsedTime;
175-
updateCellExecutionTimes(this.editor, this.cell, {
181+
updateCellExecutionTimes(this.cell, {
176182
startTime: this.cell.metadata.runStartTime,
183+
lastRunDuration: this.stopWatch.elapsedTime,
177184
duration: this.cell.metadata.lastRunDuration
178185
});
186+
179187
// If there are any errors in the cell, then change status to error.
180188
if (this.cell.outputs.some((output) => output.outputKind === vscodeNotebookEnums.CellOutputKind.Error)) {
181-
this.cell.metadata.runState = vscodeNotebookEnums.NotebookCellRunState.Error;
182-
this.cell.metadata.statusMessage = getCellStatusMessageBasedOnFirstCellErrorOutput(this.cell.outputs);
189+
runState = vscodeNotebookEnums.NotebookCellRunState.Error;
190+
statusMessage = getCellStatusMessageBasedOnFirstCellErrorOutput(this.cell.outputs);
183191
}
184192

193+
const cellIndex = this.editor.document.cells.indexOf(this.cell);
194+
new WorkspaceEdit().replaceCellMetadata(this.cell.notebook.uri, cellIndex, {
195+
...this.cell.metadata,
196+
runState,
197+
statusMessage
198+
});
199+
185200
this._completed = true;
186201
this._result.resolve(this.cell.metadata.runState);
187202
// Changes to metadata must be saved in ipynb, hence mark doc has dirty.
@@ -207,12 +222,16 @@ export class CellExecution {
207222
* At this point we revert cell state & indicate that it has nto started & it is not busy.
208223
*/
209224
private dequeue() {
210-
if (this.oldCellRunState === vscodeNotebookEnums.NotebookCellRunState.Running) {
211-
this.cell.metadata.runState = vscodeNotebookEnums.NotebookCellRunState.Idle;
212-
} else {
213-
this.cell.metadata.runState = this.oldCellRunState;
214-
}
225+
const runState =
226+
this.oldCellRunState === vscodeNotebookEnums.NotebookCellRunState.Running
227+
? vscodeNotebookEnums.NotebookCellRunState.Idle
228+
: this.oldCellRunState;
215229
this.cell.metadata.runStartTime = undefined;
230+
new WorkspaceEdit().replaceCellMetadata(this.cell.notebook.uri, this.cell.notebook.cells.indexOf(this.cell), {
231+
...this.cell.metadata,
232+
runStartTime: undefined,
233+
runState
234+
});
216235
this._completed = true;
217236
this._result.resolve(this.cell.metadata.runState);
218237
// Changes to metadata must be saved in ipynb, hence mark doc has dirty.
@@ -224,7 +243,10 @@ export class CellExecution {
224243
* (mark it as busy).
225244
*/
226245
private enqueue() {
227-
this.cell.metadata.runState = vscodeNotebookEnums.NotebookCellRunState.Running;
246+
new WorkspaceEdit().replaceCellMetadata(this.cell.notebook.uri, this.cell.notebook.cells.indexOf(this.cell), {
247+
...this.cell.metadata,
248+
runState: vscodeNotebookEnums.NotebookCellRunState.Running
249+
});
228250
this.contentProvider.notifyChangesToDocument(this.cell.notebook);
229251
}
230252

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

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import type { nbformat } from '@jupyterlab/coreutils';
77
import type { KernelMessage } from '@jupyterlab/services';
88
import * as fastDeepEqual from 'fast-deep-equal';
9+
import { WorkspaceEdit } from 'vscode';
910
import type { NotebookCell, NotebookEditor } from '../../../../../types/vscode-proposed';
1011
import { createErrorOutput } from '../../../../datascience-ui/common/cellFactory';
1112
import { createIOutputFromCellOutputs, createVSCCellOutputsFromOutputs, translateErrorOutput } from './helpers';
@@ -66,22 +67,24 @@ export function handleUpdateDisplayDataMessage(
6667
/**
6768
* Updates the VSC cell with the error output.
6869
*/
69-
export function updateCellWithErrorStatus(editor: NotebookEditor, cell: NotebookCell, ex: Partial<Error>) {
70-
editor.edit((edit) => {
71-
const cellIndex = editor.document.cells.indexOf(cell);
72-
edit.replaceMetadata(cellIndex, { ...cell.metadata, runState: vscodeNotebookEnums.NotebookCellRunState.Error });
73-
edit.replaceOutput(cellIndex, [translateErrorOutput(createErrorOutput(ex))]);
70+
export function updateCellWithErrorStatus(cell: NotebookCell, ex: Partial<Error>) {
71+
const cellIndex = cell.notebook.cells.indexOf(cell);
72+
new WorkspaceEdit().replaceCellMetadata(cell.document.uri, cellIndex, {
73+
...cell.metadata,
74+
runState: vscodeNotebookEnums.NotebookCellRunState.Error
7475
});
76+
new WorkspaceEdit().replaceCellOutput(cell.document.uri, cellIndex, [translateErrorOutput(createErrorOutput(ex))]);
7577
}
7678

7779
/**
7880
* @returns {boolean} Returns `true` if execution count has changed.
7981
*/
8082
export function updateCellExecutionCount(editor: NotebookEditor, cell: NotebookCell, executionCount: number): boolean {
8183
if (cell.metadata.executionOrder !== executionCount && executionCount) {
82-
editor.edit((edit) => {
83-
const cellIndex = editor.document.cells.indexOf(cell);
84-
edit.replaceMetadata(cellIndex, { ...cell.metadata, executionOrder: executionCount });
84+
const cellIndex = editor.document.cells.indexOf(cell);
85+
new WorkspaceEdit().replaceCellMetadata(cell.document.uri, cellIndex, {
86+
...cell.metadata,
87+
executionOrder: executionCount
8588
});
8689
return true;
8790
}
@@ -94,21 +97,18 @@ export function updateCellExecutionCount(editor: NotebookEditor, cell: NotebookC
9497
* Here we update both the VSCode Cell as well as our ICell (cell in our INotebookModel).
9598
* @returns {(boolean | undefined)} Returns `true` if output has changed.
9699
*/
97-
export function updateCellOutput(
98-
editor: NotebookEditor,
99-
vscCell: NotebookCell,
100-
outputs: nbformat.IOutput[]
101-
): boolean | undefined {
100+
export function updateCellOutput(cell: NotebookCell, outputs: nbformat.IOutput[]): boolean | undefined {
102101
const newOutput = createVSCCellOutputsFromOutputs(outputs);
103102
// If there was no output and still no output, then nothing to do.
104-
if (vscCell.outputs.length === 0 && newOutput.length === 0) {
103+
if (cell.outputs.length === 0 && newOutput.length === 0) {
105104
return;
106105
}
107106
// Compare outputs (at the end of the day everything is serializable).
108107
// Hence this is a safe comparison.
109-
if (vscCell.outputs.length === newOutput.length && fastDeepEqual(vscCell.outputs, newOutput)) {
108+
if (cell.outputs.length === newOutput.length && fastDeepEqual(cell.outputs, newOutput)) {
110109
return;
111110
}
112-
editor.edit((edit) => edit.replaceOutput(0, newOutput));
111+
const cellIndex = cell.notebook.cells.indexOf(cell);
112+
new WorkspaceEdit().replaceCellOutput(cell.document.uri, cellIndex, newOutput);
113113
return true;
114114
}

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

Lines changed: 61 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,9 @@ import { JupyterNotebookView } from '../constants';
2727
// tslint:disable-next-line: no-var-requires no-require-imports
2828
const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed');
2929
// tslint:disable-next-line: no-require-imports
30+
import { url } from 'inspector';
3031
import cloneDeep = require('lodash/cloneDeep');
32+
import { WorkspaceEdit } from 'vscode';
3133
import { isUntitledFile } from '../../../common/utils/misc';
3234
import { KernelConnectionMetadata } from '../../jupyter/kernels/types';
3335
import { updateNotebookMetadata } from '../../notebookStorage/baseModel';
@@ -339,58 +341,59 @@ export function createIOutputFromCellOutputs(cellOutputs: CellOutput[]): nbforma
339341
.map((output) => output!!);
340342
}
341343

342-
export function clearCellForExecution(editor: NotebookEditor, cell: NotebookCell) {
343-
editor.edit((edit) => {
344-
const cellIndex = editor.document.cells.indexOf(cell);
345-
edit.replaceMetadata(cellIndex, {
346-
...cell.metadata,
347-
statusMessage: undefined,
348-
executionOrder: undefined,
349-
lastRunDuration: undefined,
350-
runStartTime: undefined
351-
});
352-
edit.replaceOutput(cellIndex, []);
344+
export function clearCellForExecution(cell: NotebookCell) {
345+
const cellIndex = cell.notebook.cells.indexOf(cell);
346+
new WorkspaceEdit().replaceCellMetadata(cell.notebook.uri, cellIndex, {
347+
...cell.metadata,
348+
statusMessage: undefined,
349+
executionOrder: undefined,
350+
lastRunDuration: undefined,
351+
runStartTime: undefined
353352
});
353+
new WorkspaceEdit().replaceCellOutput(cell.notebook.uri, cellIndex, []);
354354

355-
updateCellExecutionTimes(editor, cell);
355+
updateCellExecutionTimes(cell);
356356
}
357357

358358
/**
359359
* Store execution start and end times.
360360
* Stored as ISO for portability.
361361
*/
362362
export function updateCellExecutionTimes(
363-
editor: NotebookEditor,
364363
cell: NotebookCell,
365-
times?: { startTime?: number; duration?: number }
364+
times?: { startTime?: number; duration?: number; lastRunDuration?: number }
366365
) {
367-
editor.edit((edit) => {
368-
const cellIndex = editor.document.cells.indexOf(cell);
369-
if (!times || !times.duration || !times.startTime) {
370-
const cellMetadata = cloneDeep(cell.metadata);
371-
let updated = false;
372-
if (cellMetadata.custom?.metadata?.vscode?.start_execution_time) {
373-
delete cellMetadata.custom.metadata.vscode.start_execution_time;
374-
updated = true;
375-
}
376-
if (cellMetadata.custom?.metadata?.vscode?.end_execution_time) {
377-
delete cellMetadata.custom.metadata.vscode.end_execution_time;
378-
updated = true;
379-
}
380-
if (updated) {
381-
edit.replaceMetadata(cellIndex, { ...cellMetadata });
382-
}
383-
return;
366+
const cellIndex = cell.notebook.cells.indexOf(cell);
367+
368+
if (!times || !times.duration || !times.startTime) {
369+
const cellMetadata = cloneDeep(cell.metadata);
370+
let updated = false;
371+
if (cellMetadata.custom?.metadata?.vscode?.start_execution_time) {
372+
delete cellMetadata.custom.metadata.vscode.start_execution_time;
373+
updated = true;
374+
}
375+
if (cellMetadata.custom?.metadata?.vscode?.end_execution_time) {
376+
delete cellMetadata.custom.metadata.vscode.end_execution_time;
377+
updated = true;
384378
}
379+
if (updated) {
380+
new WorkspaceEdit().replaceCellMetadata(cell.notebook.uri, cellIndex, { ...cellMetadata });
381+
}
382+
return;
383+
}
385384

386-
const startTimeISO = new Date(times.startTime).toISOString();
387-
const endTimeISO = new Date(times.startTime + times.duration).toISOString();
388-
const customMetadata = cloneDeep(cell.metadata.custom || {});
389-
customMetadata.metadata = customMetadata.metadata || {};
390-
customMetadata.metadata.vscode = customMetadata.metadata.vscode || {};
391-
customMetadata.metadata.vscode.end_execution_time = endTimeISO;
392-
customMetadata.metadata.vscode.start_execution_time = startTimeISO;
393-
edit.replaceMetadata(cellIndex, { ...cell.metadata, custom: customMetadata });
385+
const startTimeISO = new Date(times.startTime).toISOString();
386+
const endTimeISO = new Date(times.startTime + times.duration).toISOString();
387+
const customMetadata = cloneDeep(cell.metadata.custom || {});
388+
customMetadata.metadata = customMetadata.metadata || {};
389+
customMetadata.metadata.vscode = customMetadata.metadata.vscode || {};
390+
customMetadata.metadata.vscode.end_execution_time = endTimeISO;
391+
customMetadata.metadata.vscode.start_execution_time = startTimeISO;
392+
const lastRunDuration = times.lastRunDuration ?? cell.metadata.lastRunDuration;
393+
new WorkspaceEdit().replaceCellMetadata(cell.notebook.uri, cellIndex, {
394+
...cell.metadata,
395+
custom: customMetadata,
396+
lastRunDuration
394397
});
395398
}
396399

@@ -662,11 +665,7 @@ export function getCellStatusMessageBasedOnFirstCellErrorOutput(outputs?: CellOu
662665
/**
663666
* Updates a notebook document as a result of trusting it.
664667
*/
665-
export function updateVSCNotebookAfterTrustingNotebook(
666-
editor: NotebookEditor,
667-
document: NotebookDocument,
668-
originalCells: ICell[]
669-
) {
668+
export function updateVSCNotebookAfterTrustingNotebook(document: NotebookDocument, originalCells: ICell[]) {
670669
const areAllCellsEditableAndRunnable = document.cells.every((cell) => {
671670
if (cell.cellKind === vscodeNotebookEnums.CellKind.Markdown) {
672671
return cell.metadata.editable;
@@ -690,16 +689,23 @@ export function updateVSCNotebookAfterTrustingNotebook(
690689
document.metadata.editable = true;
691690
document.metadata.runnable = true;
692691

693-
editor.edit((edit) => {
694-
document.cells.forEach((cell, index) => {
695-
if (cell.cellKind === vscodeNotebookEnums.CellKind.Markdown) {
696-
edit.replaceMetadata(index, { ...cell.metadata, editable: true });
697-
} else {
698-
edit.replaceMetadata(index, { ...cell.metadata, editable: true, runnable: true });
699-
// Restore the output once we trust the notebook.
700-
// tslint:disable-next-line: no-any
701-
edit.replaceOutput(index, createVSCCellOutputsFromOutputs(originalCells[index].data.outputs as any));
702-
}
703-
});
692+
const workspaceEdit = new WorkspaceEdit();
693+
document.cells.forEach((cell, index) => {
694+
if (cell.cellKind === vscodeNotebookEnums.CellKind.Markdown) {
695+
workspaceEdit.replaceCellMetadata(document.uri, index, { ...cell.metadata, editable: true });
696+
} else {
697+
workspaceEdit.replaceCellMetadata(document.uri, index, {
698+
...cell.metadata,
699+
editable: true,
700+
runnable: true
701+
});
702+
// Restore the output once we trust the notebook.
703+
// tslint:disable-next-line: no-any
704+
workspaceEdit.replaceCellOutput(
705+
document.uri,
706+
index,
707+
createVSCCellOutputsFromOutputs(originalCells[index].data.outputs as any)
708+
);
709+
}
704710
});
705711
}

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

Lines changed: 21 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
'use strict';
55

6-
import { ConfigurationTarget, Event, EventEmitter, Uri, WebviewPanel } from 'vscode';
6+
import { ConfigurationTarget, Event, EventEmitter, Uri, WebviewPanel, WorkspaceEdit } from 'vscode';
77
import type { NotebookDocument } from 'vscode-proposed';
88
import { IApplicationShell, ICommandManager, IVSCodeNotebook } from '../../common/application/types';
99
import { traceError } from '../../common/logger';
@@ -128,40 +128,39 @@ export class NotebookEditor implements INotebookEditor {
128128
return;
129129
}
130130
const defaultLanguage = getDefaultCodeLanguage(this.model);
131-
this.vscodeNotebook.activeNotebookEditor.edit((editor) => {
132-
const totalLength = this.document.cells.length;
133-
editor.insert(
134-
this.document.cells.length,
135-
'',
136-
defaultLanguage,
137-
vscodeNotebookEnums.CellKind.Code,
138-
[],
139-
undefined
140-
);
141-
for (let i = totalLength - 1; i >= 0; i = i - 1) {
142-
editor.delete(i);
131+
new WorkspaceEdit().replaceCells(this.document.uri, 0, 0, [
132+
{
133+
cellKind: vscodeNotebookEnums.CellKind.Code,
134+
language: defaultLanguage,
135+
metadata: {},
136+
outputs: [],
137+
source: ''
143138
}
144-
});
139+
]);
145140
}
146141
public expandAllCells(): void {
147142
if (!this.vscodeNotebook.activeNotebookEditor) {
148143
return;
149144
}
150-
const cells = this.vscodeNotebook.activeNotebookEditor.document.cells;
151-
this.vscodeNotebook.activeNotebookEditor.edit((edit) => {
152-
cells.forEach((cell, index) => {
153-
edit.replaceMetadata(index, { ...cell.metadata, inputCollapsed: false, outputCollapsed: false });
145+
const notebook = this.vscodeNotebook.activeNotebookEditor.document;
146+
notebook.cells.forEach((cell, index) => {
147+
new WorkspaceEdit().replaceCellMetadata(notebook.uri, index, {
148+
...cell.metadata,
149+
inputCollapsed: false,
150+
outputCollapsed: false
154151
});
155152
});
156153
}
157154
public collapseAllCells(): void {
158155
if (!this.vscodeNotebook.activeNotebookEditor) {
159156
return;
160157
}
161-
const cells = this.vscodeNotebook.activeNotebookEditor.document.cells;
162-
this.vscodeNotebook.activeNotebookEditor.edit((edit) => {
163-
cells.forEach((cell, index) => {
164-
edit.replaceMetadata(index, { ...cell.metadata, inputCollapsed: true, outputCollapsed: true });
158+
const notebook = this.vscodeNotebook.activeNotebookEditor.document;
159+
notebook.cells.forEach((cell, index) => {
160+
new WorkspaceEdit().replaceCellMetadata(notebook.uri, index, {
161+
...cell.metadata,
162+
inputCollapsed: true,
163+
outputCollapsed: true
165164
});
166165
});
167166
}

0 commit comments

Comments
 (0)