Check notebook URIs in BulkCellEdits

src/vs/workbench/api/common/extHostTypes.ts

@@ -17,7 +17,7 @@ import { IExtensionDescription } from 'vs/platform/extensions/common/extensions'
 import { FileSystemProviderErrorCode, markAsFileSystemProviderError } from 'vs/platform/files/common/files';
 import { RemoteAuthorityResolverErrorCode } from 'vs/platform/remote/common/remoteAuthorityResolver';
 import { IRelativePatternDto } from 'vs/workbench/api/common/extHost.protocol';
-import { CellEditType, CellUri, ICellPartialMetadataEdit, IDocumentMetadataEdit, isTextStreamMime } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellEditType, ICellPartialMetadataEdit, IDocumentMetadataEdit, isTextStreamMime } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { checkProposedApiEnabled } from 'vs/workbench/services/extensions/common/extensions';
 import type * as vscode from 'vscode';
 
@@ -862,10 +862,6 @@ export class WorkspaceEdit implements vscode.WorkspaceEdit {
 					edit = editOrTuple;
 				}
 				if (NotebookEdit.isNotebookCellEdit(edit)) {
-					if (uri.scheme === CellUri.scheme) {
-						throw new Error('set must be called with a notebook document URI, not a cell URI.');
-					}
-
 					if (edit.newCellMetadata) {
 						this.replaceNotebookCellMetadata(uri, edit.range.start, edit.newCellMetadata, metadata);
 					} else if (edit.newNotebookMetadata) {

src/vs/workbench/api/test/browser/extHostTypes.test.ts

@@ -4,14 +4,13 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as assert from 'assert';
-import { CancellationError } from 'vs/base/common/errors';
-import { MarshalledId } from 'vs/base/common/marshallingIds';
-import { Mimes } from 'vs/base/common/mime';
-import { isWindows } from 'vs/base/common/platform';
-import { assertType } from 'vs/base/common/types';
 import { URI } from 'vs/base/common/uri';
 import * as types from 'vs/workbench/api/common/extHostTypes';
-import { CellUri } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { isWindows } from 'vs/base/common/platform';
+import { assertType } from 'vs/base/common/types';
+import { Mimes } from 'vs/base/common/mime';
+import { MarshalledId } from 'vs/base/common/marshallingIds';
+import { CancellationError } from 'vs/base/common/errors';
 
 function assertToJSON(a: any, expected: any) {
 	const raw = JSON.stringify(a);
@@ -433,23 +432,6 @@ suite('ExtHostTypes', function () {
 		assert.strictEqual(second.edit.newText, 'Foo');
 	});
 
-	test('WorkspaceEdit - NotebookEdits', () => {
-		const edit = new types.WorkspaceEdit();
-		const notebookEdit = types.NotebookEdit.insertCells(0, [new types.NotebookCellData(types.NotebookCellKind.Code, '// hello', 'javascript')]) as types.NotebookEdit;
-		const notebookUri = URI.parse('/foo/notebook.ipynb');
-		edit.set(notebookUri, [notebookEdit]);
-
-		const cellUri = CellUri.generate(notebookUri, 123);
-		try {
-			edit.set(cellUri, [notebookEdit]);
-		} catch (err) {
-			assert.ok(err.message.includes('set must be called with a notebook document URI'), err.toString());
-			return;
-		}
-
-		throw new Error('Expected set to throw with cell URI');
-	});
-
 	test('DocumentLink', () => {
 		assert.throws(() => new types.DocumentLink(null!, null!));
 		assert.throws(() => new types.DocumentLink(new types.Range(1, 1, 1, 1), null!));

src/vs/workbench/contrib/bulkEdit/browser/bulkCellEdits.ts

@@ -13,7 +13,7 @@ import { WorkspaceEditMetadata } from 'vs/editor/common/languages';
 import { IProgress } from 'vs/platform/progress/common/progress';
 import { UndoRedoGroup, UndoRedoSource } from 'vs/platform/undoRedo/common/undoRedo';
 import { getNotebookEditorFromEditorPane } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
-import { ICellPartialMetadataEdit, ICellReplaceEdit, IDocumentMetadataEdit, ISelectionState, IWorkspaceNotebookCellEdit, SelectionStateType } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellUri, ICellPartialMetadataEdit, ICellReplaceEdit, IDocumentMetadataEdit, ISelectionState, IWorkspaceNotebookCellEdit, SelectionStateType } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { INotebookEditorModelResolverService } from 'vs/workbench/contrib/notebook/common/notebookEditorModelResolverService';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 
@@ -54,7 +54,20 @@ export class BulkCellEdits {
 		private readonly _edits: ResourceNotebookCellEdit[],
 		@IEditorService private readonly _editorService: IEditorService,
 		@INotebookEditorModelResolverService private readonly _notebookModelService: INotebookEditorModelResolverService,
-	) { }
+	) {
+		this._edits = this._edits.map(e => {
+			if (e.resource.scheme === CellUri.scheme) {
+				const uri = CellUri.parse(e.resource)?.notebook;
+				if (!uri) {
+					throw new Error(`Invalid notebook URI: ${e.resource}`);
+				}
+
+				return new ResourceNotebookCellEdit(uri, e.cellEdit, e.notebookVersionId, e.metadata);
+			} else {
+				return e;
+			}
+		});
+	}
 
 	async apply(): Promise<readonly URI[]> {
 		const resources: URI[] = [];

src/vs/workbench/contrib/bulkEdit/test/browser/bulkCellEdits.test.ts

@@ -0,0 +1,55 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as assert from 'assert';
+import { CancellationTokenSource } from 'vs/base/common/cancellation';
+import { URI } from 'vs/base/common/uri';
+import { mockObject } from 'vs/base/test/common/mock';
+import { IProgress } from 'vs/platform/progress/common/progress';
+import { UndoRedoGroup, UndoRedoSource } from 'vs/platform/undoRedo/common/undoRedo';
+import { BulkCellEdits, ResourceNotebookCellEdit } from 'vs/workbench/contrib/bulkEdit/browser/bulkCellEdits';
+import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
+import { CellEditType, CellUri, IResolvedNotebookEditorModel } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookEditorModelResolverService } from 'vs/workbench/contrib/notebook/common/notebookEditorModelResolverService';
+import { TestEditorService } from 'vs/workbench/test/browser/workbenchTestServices';
+
+suite('BulkCellEdits', function () {
+	async function runTest(inputUri: URI, resolveUri: URI) {
+		const progress: IProgress<void> = { report: _ => { } };
+		const editorService = new TestEditorService();
+
+		const notebook = mockObject<NotebookTextModel>()();
+		notebook.uri.returns(URI.file('/project/notebook.ipynb'));
+
+		const notebookEditorModel = mockObject<IResolvedNotebookEditorModel>()({ notebook: notebook as any });
+		notebookEditorModel.isReadonly.returns(false);
+
+		const notebookService = mockObject<INotebookEditorModelResolverService>()();
+		notebookService.resolve.returns({ object: notebookEditorModel, dispose: () => { } });
+
+		const edits = [
+			new ResourceNotebookCellEdit(inputUri, { index: 0, count: 1, editType: CellEditType.Replace, cells: [] })
+		];
+		const bce = new BulkCellEdits(new UndoRedoGroup(), new UndoRedoSource(), progress, new CancellationTokenSource().token, edits, editorService, notebookService as any);
+		await bce.apply();
+
+		const resolveArgs = notebookService.resolve.args[0];
+		assert.strictEqual(resolveArgs[0].toString(), resolveUri.toString());
+	}
+
+	const notebookUri = URI.file('/foo/bar.ipynb');
+	test('works with notebook URI', async () => {
+		await runTest(notebookUri, notebookUri);
+	});
+
+	test('maps cell URI to notebook URI', async () => {
+		await runTest(CellUri.generate(notebookUri, 5), notebookUri);
+	});
+
+	test('throws for invalid cell URI', async () => {
+		const badCellUri = CellUri.generate(notebookUri, 5).with({ fragment: '' });
+		await assert.rejects(async () => await runTest(badCellUri, notebookUri));
+	});
+});

test/unit/assert.js

@@ -467,5 +467,52 @@ assert.doesNotThrow = function(block, /*optional*/message) {
 };
 
 assert.ifError = function(err) { if (err) {throw err;}};
+
+function checkIsPromise(obj) {
+	return (obj !== null && typeof obj === 'object' &&
+		typeof obj.then === 'function' &&
+		typeof obj.catch === 'function');
+}
+
+const NO_EXCEPTION_SENTINEL = {};
+async function waitForActual(promiseFn) {
+	let resultPromise;
+	if (typeof promiseFn === 'function') {
+		// Return a rejected promise if `promiseFn` throws synchronously.
+		resultPromise = promiseFn();
+		// Fail in case no promise is returned.
+		if (!checkIsPromise(resultPromise)) {
+			throw new Error('ERR_INVALID_RETURN_VALUE: promiseFn did not return Promise. ' + resultPromise);
+		}
+	} else if (checkIsPromise(promiseFn)) {
+		resultPromise = promiseFn;
+	} else {
+		throw new Error('ERR_INVALID_ARG_TYPE: promiseFn is not Function or Promise. ' + promiseFn);
+	}
+
+	try {
+		await resultPromise;
+	} catch (e) {
+		return e;
+	}
+	return NO_EXCEPTION_SENTINEL;
+}
+
+function expectsError(shouldHaveError, actual, message) {
+	if (shouldHaveError && actual === NO_EXCEPTION_SENTINEL) {
+		fail(undefined, 'Error', `Missing expected rejection${message ? ': ' + message : ''}`)
+	} else if (!shouldHaveError && actual !== NO_EXCEPTION_SENTINEL) {
+		fail(actual, undefined, `Got unexpected rejection (${actual.message})${message ? ': ' + message : ''}`)
+	}
+}
+
+assert.rejects = async function rejects(promiseFn, message) {
+	expectsError(true, await waitForActual(promiseFn), message);
+};
+
+assert.doesNotReject = async function doesNotReject(fn, message) {
+	expectsError(false, await waitForActual(fn), message);
+};
+
 return assert;
 });