Add notebook Format On Cell Execution

src/vs/workbench/contrib/notebook/browser/contrib/format/formatting.ts

@@ -6,7 +6,7 @@
 import { localize } from 'vs/nls';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { KeyCode, KeyMod } from 'vs/base/common/keyCodes';
-import { DisposableStore } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableStore } from 'vs/base/common/lifecycle';
 import { ICodeEditor } from 'vs/editor/browser/editorBrowser';
 import { EditorAction, registerEditorAction } from 'vs/editor/browser/editorExtensions';
 import { IBulkEditService, ResourceTextEdit } from 'vs/editor/browser/services/bulkEditService';
@@ -24,6 +24,14 @@ import { NOTEBOOK_ACTIONS_CATEGORY } from 'vs/workbench/contrib/notebook/browser
 import { getNotebookEditorFromEditorPane } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { NOTEBOOK_EDITOR_EDITABLE, NOTEBOOK_IS_ACTIVE_EDITOR } from 'vs/workbench/contrib/notebook/common/notebookContextKeys';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
+import { INotebookCellExecution } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
+import { ICellExecutionParticipant, INotebookExecutionService } from 'vs/workbench/contrib/notebook/common/notebookExecutionService';
+import { NotebookSetting } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { LifecyclePhase } from 'vs/workbench/services/lifecycle/common/lifecycle';
+import { Registry } from 'vs/platform/registry/common/platform';
+import { IWorkbenchContribution, IWorkbenchContributionsRegistry, Extensions as WorkbenchContributionsExtensions } from 'vs/workbench/common/contributions';
+import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
 
 // format notebook
 registerAction2(class extends Action2 {
@@ -74,7 +82,8 @@ registerAction2(class extends Action2 {
 					editorWorkerService,
 					languageFeaturesService,
 					model,
-					model.getOptions(), CancellationToken.None
+					model.getOptions(),
+					CancellationToken.None
 				);
 
 				const edits: ResourceTextEdit[] = [];
@@ -126,3 +135,88 @@ registerEditorAction(class FormatCellAction extends EditorAction {
 		}
 	}
 });
+
+class FormatOnCellExecutionParticipant implements ICellExecutionParticipant {
+	constructor(
+		@IBulkEditService private readonly bulkEditService: IBulkEditService,
+		@ILanguageFeaturesService private readonly languageFeaturesService: ILanguageFeaturesService,
+		@ITextModelService private readonly textModelService: ITextModelService,
+		@IEditorWorkerService private readonly editorWorkerService: IEditorWorkerService,
+		@IConfigurationService private readonly configurationService: IConfigurationService,
+		@INotebookService private readonly _notebookService: INotebookService,
+	) {
+	}
+
+	async onWillExecuteCell(executions: INotebookCellExecution[]): Promise<void> {
+
+		const enabled = this.configurationService.getValue<boolean>(NotebookSetting.formatOnCellExecution);
+		if (!enabled) {
+			return;
+		}
+
+		const disposable = new DisposableStore();
+		try {
+			const allCellEdits = await Promise.all(executions.map(async cellExecution => {
+				const nbModel = this._notebookService.getNotebookTextModel(cellExecution.notebook);
+				if (!nbModel) {
+					return [];
+				}
+				let activeCell;
+				for (const cell of nbModel.cells) {
+					if (cell.handle === cellExecution.cellHandle) {
+						activeCell = cell;
+						break;
+					}
+				}
+				if (!activeCell) {
+					return [];
+				}
+
+				const ref = await this.textModelService.createModelReference(activeCell.uri);
+				disposable.add(ref);
+
+				const model = ref.object.textEditorModel;
+
+				// todo: eventually support cancellation. potential leak if cell deleted mid execution
+				const formatEdits = await getDocumentFormattingEditsUntilResult(
+					this.editorWorkerService,
+					this.languageFeaturesService,
+					model,
+					model.getOptions(),
+					CancellationToken.None
+				);
+
+				const edits: ResourceTextEdit[] = [];
+
+				if (formatEdits) {
+					edits.push(...formatEdits.map(edit => new ResourceTextEdit(model.uri, edit, model.getVersionId())));
+					return edits;
+				}
+
+				return [];
+			}));
+
+			await this.bulkEditService.apply(/* edit */allCellEdits.flat(), { label: localize('formatCells.label', "Format Cells"), code: 'undoredo.notebooks.onWillExecuteFormat', });
+
+		} finally {
+			disposable.dispose();
+		}
+	}
+}
+
+export class CellExecutionParticipantsContribution extends Disposable implements IWorkbenchContribution {
+	constructor(
+		@IInstantiationService private readonly instantiationService: IInstantiationService,
+		@INotebookExecutionService private readonly notebookExecutionService: INotebookExecutionService
+	) {
+		super();
+		this.registerKernelExecutionParticipants();
+	}
+
+	private registerKernelExecutionParticipants(): void {
+		this._register(this.notebookExecutionService.registerExecutionParticipant(this.instantiationService.createInstance(FormatOnCellExecutionParticipant)));
+	}
+}
+
+const workbenchContributionsRegistry = Registry.as<IWorkbenchContributionsRegistry>(WorkbenchContributionsExtensions.Workbench);
+workbenchContributionsRegistry.registerWorkbenchContribution(CellExecutionParticipantsContribution, LifecyclePhase.Restored);

src/vs/workbench/contrib/notebook/browser/notebook.contribution.ts

@@ -926,9 +926,13 @@ configurationRegistry.registerConfiguration({
 			additionalProperties: {
 				type: 'boolean'
 			},
-			tags: ['notebookLayout'],
 			default: {}
 		},
+		[NotebookSetting.formatOnCellExecution]: {
+			markdownDescription: nls.localize('notebook.formatOnCellExecution', "Format a notebook cell upon execution. A formatter must be available."),
+			type: 'boolean',
+			default: false
+		},
 		[NotebookSetting.confirmDeleteRunningCell]: {
 			markdownDescription: nls.localize('notebook.confirmDeleteRunningCell', "Control whether a confirmation prompt is required to delete a running cell."),
 			type: 'boolean',

src/vs/workbench/contrib/notebook/browser/services/notebookExecutionServiceImpl.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { CancellationTokenSource } from 'vs/base/common/cancellation';
-import { IDisposable } from 'vs/base/common/lifecycle';
+import { IDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import * as nls from 'vs/nls';
 import { ICommandService } from 'vs/platform/commands/common/commands';
 import { IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
@@ -13,10 +13,11 @@ import { IWorkspaceTrustRequestService } from 'vs/platform/workspace/common/work
 import { KernelPickerMRUStrategy } from 'vs/workbench/contrib/notebook/browser/viewParts/notebookKernelQuickPickStrategy';
 import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
 import { CellKind, INotebookTextModel, NotebookCellExecutionState } from 'vs/workbench/contrib/notebook/common/notebookCommon';
-import { INotebookExecutionService } from 'vs/workbench/contrib/notebook/common/notebookExecutionService';
+import { INotebookExecutionService, ICellExecutionParticipant } from 'vs/workbench/contrib/notebook/common/notebookExecutionService';
 import { INotebookCellExecution, INotebookExecutionStateService } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
 import { INotebookKernelHistoryService, INotebookKernelService } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
 
+
 export class NotebookExecutionService implements INotebookExecutionService, IDisposable {
 	declare _serviceBrand: undefined;
 	private _activeProxyKernelExecutionToken: CancellationTokenSource | undefined;
@@ -77,6 +78,8 @@ export class NotebookExecutionService implements INotebookExecutionService, IDis
 
 		// request execution
 		if (validCellExecutions.length > 0) {
+			await this.runExecutionParticipants(validCellExecutions);
+
 			this._notebookKernelService.selectKernelForNotebook(kernel, notebook);
 			await kernel.executeNotebookCellsRequest(notebook.uri, validCellExecutions.map(c => c.cellHandle));
 			// the connecting state can change before the kernel resolves executeNotebookCellsRequest
@@ -102,6 +105,20 @@ export class NotebookExecutionService implements INotebookExecutionService, IDis
 		this.cancelNotebookCellHandles(notebook, Array.from(cells, cell => cell.handle));
 	}
 
+	private readonly cellExecutionParticipants = new Set<ICellExecutionParticipant>;
+
+	registerExecutionParticipant(participant: ICellExecutionParticipant) {
+		this.cellExecutionParticipants.add(participant);
+		return toDisposable(() => this.cellExecutionParticipants.delete(participant));
+	}
+
+	async runExecutionParticipants(executions: INotebookCellExecution[]): Promise<void> {
+		for (const participant of this.cellExecutionParticipants) {
+			await participant.onWillExecuteCell(executions);
+		}
+		return;
+	}
+
 	dispose() {
 		this._activeProxyKernelExecutionToken?.dispose(true);
 	}

src/vs/workbench/contrib/notebook/common/notebookCommon.ts

@@ -938,6 +938,7 @@ export const NotebookSetting = {
 	outputScrolling: 'notebook.output.scrolling',
 	textOutputLineLimit: 'notebook.output.textLineLimit',
 	formatOnSave: 'notebook.formatOnSave.enabled',
+	formatOnCellExecution: 'notebook.formatOnCellExecution',
 	codeActionsOnSave: 'notebook.codeActionsOnSave',
 	outputWordWrap: 'notebook.output.wordWrap',
 	outputLineHeightDeprecated: 'notebook.outputLineHeight',

src/vs/workbench/contrib/notebook/common/notebookExecutionService.ts

@@ -3,10 +3,12 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { IDisposable } from 'vs/base/common/lifecycle';
 import { IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
 import { INotebookTextModel, IOutputDto, IOutputItemDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookCellExecution } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
 
 export enum CellExecutionUpdateType {
 	Output = 1,
@@ -36,4 +38,10 @@ export interface INotebookExecutionService {
 	executeNotebookCells(notebook: INotebookTextModel, cells: Iterable<NotebookCellTextModel>, contextKeyService: IContextKeyService): Promise<void>;
 	cancelNotebookCells(notebook: INotebookTextModel, cells: Iterable<NotebookCellTextModel>): Promise<void>;
 	cancelNotebookCellHandles(notebook: INotebookTextModel, cells: Iterable<number>): Promise<void>;
+	registerExecutionParticipant(participant: ICellExecutionParticipant): IDisposable;
+	runExecutionParticipants(executions: INotebookCellExecution[]): Promise<void>;
+}
+
+export interface ICellExecutionParticipant {
+	onWillExecuteCell(executions: INotebookCellExecution[]): Promise<void>;
 }