Introduce notebook contribution.

src/vs/workbench/contrib/notebook/browser/contrib/fold/folding.ts

@@ -3,14 +3,26 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Disposable } from 'vs/base/common/lifecycle';
-import { INotebookEditor, INotebookEditorMouseEvent, ICellRange } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
+import { Disposable, DisposableStore } from 'vs/base/common/lifecycle';
+import { INotebookEditor, INotebookEditorMouseEvent, ICellRange, INotebookEditorContribution, NOTEBOOK_EDITOR_FOCUSED } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import * as DOM from 'vs/base/browser/dom';
 import { CellFoldingState, FoldingModel } from 'vs/workbench/contrib/notebook/browser/contrib/fold/foldingModel';
 import { CellKind } from 'vs/workbench/contrib/notebook/common/notebookCommon';
-
-export class FoldingController extends Disposable {
-	private _foldingModel: FoldingModel;
+import { registerNotebookContribution } from 'vs/workbench/contrib/notebook/browser/notebookEditorExtensions';
+import { registerAction2, Action2 } from 'vs/platform/actions/common/actions';
+import { ContextKeyExpr } from 'vs/platform/contextkey/common/contextkey';
+import { InputFocusedContextKey } from 'vs/platform/contextkey/common/contextkeys';
+import { KeyCode } from 'vs/base/common/keyCodes';
+import { KeybindingWeight } from 'vs/platform/keybinding/common/keybindingsRegistry';
+import { ServicesAccessor } from 'vs/platform/instantiation/common/instantiation';
+import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
+import { getActiveNotebookEditor } from 'vs/workbench/contrib/notebook/browser/contrib/notebookActions';
+
+export class FoldingController extends Disposable implements INotebookEditorContribution {
+	static id: string = 'workbench.notebook.findController';
+
+	private _foldingModel: FoldingModel | null = null;
+	private _localStore: DisposableStore = new DisposableStore();
 	constructor(
 		private readonly _notebookEditor: INotebookEditor
 
@@ -18,31 +30,45 @@ export class FoldingController extends Disposable {
 		super();
 
 		this._register(this._notebookEditor.onMouseUp(e => { this.onMouseUp(e); }));
-		this._register(this._notebookEditor.viewModel!.eventDispatcher.onDidChangeCellState(e => {
-			if (e.source.editStateChanged && e.cell.cellKind === CellKind.Markdown) {
-				this._foldingModel.recompute();
-				// this._updateEditorFoldingRanges();
+
+		this._register(this._notebookEditor.onDidChangeModel(() => {
+			this._localStore.clear();
+
+			if (!this._notebookEditor.viewModel) {
+				return;
 			}
-		}));
 
-		this._foldingModel = new FoldingModel();
-		this._foldingModel.attachViewModel(this._notebookEditor.viewModel!);
+			this._localStore.add(this._notebookEditor.viewModel!.eventDispatcher.onDidChangeCellState(e => {
+				if (e.source.editStateChanged && e.cell.cellKind === CellKind.Markdown) {
+					this._foldingModel?.recompute();
+					// this._updateEditorFoldingRanges();
+				}
+			}));
+
+			this._foldingModel = new FoldingModel();
+			this._localStore.add(this._foldingModel);
+			this._foldingModel.attachViewModel(this._notebookEditor.viewModel!);
 
-		this._register(this._foldingModel.onDidFoldingRegionChanged(() => {
-			this._updateEditorFoldingRanges();
+			this._localStore.add(this._foldingModel.onDidFoldingRegionChanged(() => {
+				this._updateEditorFoldingRanges();
+			}));
 		}));
 	}
 
-	applyMemento(state: ICellRange[]) {
-		this._foldingModel.applyMemento(state);
-		this._updateEditorFoldingRanges();
+	saveViewState(): any {
+		return this._foldingModel?.getMemento() || [];
 	}
 
-	getMemento(): ICellRange[] {
-		return this._foldingModel.getMemento();
+	restoreViewState(state: ICellRange[] | undefined) {
+		this._foldingModel?.applyMemento(state || []);
+		this._updateEditorFoldingRanges();
 	}
 
 	setFoldingState(index: number, state: CellFoldingState) {
+		if (!this._foldingModel) {
+			return;
+		}
+
 		const range = this._foldingModel.regions.findRange(index + 1);
 		const startIndex = this._foldingModel.regions.getStartLineNumber(range) - 1;
 
@@ -55,6 +81,10 @@ export class FoldingController extends Disposable {
 	}
 
 	private _updateEditorFoldingRanges() {
+		if (!this._foldingModel) {
+			return;
+		}
+
 		this._notebookEditor.viewModel!.updateFoldingRanges(this._foldingModel.regions);
 		const hiddenRanges = this._notebookEditor.viewModel!.getHiddenRanges();
 		this._notebookEditor.setHiddenAreas(hiddenRanges);
@@ -96,3 +126,77 @@ export class FoldingController extends Disposable {
 		return;
 	}
 }
+
+registerNotebookContribution(FoldingController.id, FoldingController);
+
+registerAction2(class extends Action2 {
+	constructor() {
+		super({
+			id: 'workbench.action.notebook.fold',
+			title: 'Notebook Fold Cell',
+			keybinding: {
+				when: ContextKeyExpr.and(NOTEBOOK_EDITOR_FOCUSED, ContextKeyExpr.not(InputFocusedContextKey)),
+				primary: KeyCode.LeftArrow,
+				weight: KeybindingWeight.WorkbenchContrib
+			}
+		});
+	}
+
+	async run(accessor: ServicesAccessor): Promise<void> {
+		const editorService = accessor.get(IEditorService);
+
+		const editor = getActiveNotebookEditor(editorService);
+		if (!editor) {
+			return;
+		}
+
+		const activeCell = editor.getActiveCell();
+		if (!activeCell) {
+			return;
+		}
+
+		const controller = editor.getContribution<FoldingController>(FoldingController.id);
+
+		const index = editor.viewModel?.viewCells.indexOf(activeCell);
+
+		if (index !== undefined) {
+			controller.setFoldingState(index, CellFoldingState.Collapsed);
+		}
+	}
+});
+
+registerAction2(class extends Action2 {
+	constructor() {
+		super({
+			id: 'workbench.action.notebook.unfold',
+			title: 'Notebook Unfold Cell',
+			keybinding: {
+				when: ContextKeyExpr.and(NOTEBOOK_EDITOR_FOCUSED, ContextKeyExpr.not(InputFocusedContextKey)),
+				primary: KeyCode.RightArrow,
+				weight: KeybindingWeight.WorkbenchContrib
+			}
+		});
+	}
+
+	async run(accessor: ServicesAccessor): Promise<void> {
+		const editorService = accessor.get(IEditorService);
+
+		const editor = getActiveNotebookEditor(editorService);
+		if (!editor) {
+			return;
+		}
+
+		const activeCell = editor.getActiveCell();
+		if (!activeCell) {
+			return;
+		}
+
+		const controller = editor.getContribution<FoldingController>(FoldingController.id);
+
+		const index = editor.viewModel?.viewCells.indexOf(activeCell);
+
+		if (index !== undefined) {
+			controller.setFoldingState(index, CellFoldingState.Expanded);
+		}
+	}
+});

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

@@ -30,6 +30,10 @@ import { parse } from 'vs/base/common/marshalling';
 import { CellUri, CellKind } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { ResourceMap } from 'vs/base/common/map';
 
+// Editor Contribution
+
+import 'vs/workbench/contrib/notebook/browser/contrib/fold/folding';
+
 // Output renderers registration
 
 import 'vs/workbench/contrib/notebook/browser/view/output/transforms/streamTransform';

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

@@ -96,6 +96,21 @@ export interface INotebookEditorMouseEvent {
 	readonly target: CellViewModel;
 }
 
+export interface INotebookEditorContribution {
+	/**
+	 * Dispose this contribution.
+	 */
+	dispose(): void;
+	/**
+	 * Store view state.
+	 */
+	saveViewState?(): any;
+	/**
+	 * Restore view state.
+	 */
+	restoreViewState?(state: any): void;
+}
+
 export interface INotebookEditor {
 
 	/**
@@ -103,6 +118,11 @@ export interface INotebookEditor {
 	 */
 	viewModel: NotebookViewModel | undefined;
 
+	/**
+	 * An event emitted when the model of this editor has changed.
+	 * @event
+	 */
+	readonly onDidChangeModel: Event<void>;
 	isNotebookEditor: boolean;
 
 	getInnerWebview(): Webview | undefined;
@@ -300,6 +320,13 @@ export interface INotebookEditor {
 	 * @event
 	 */
 	onMouseDown(listener: (e: INotebookEditorMouseEvent) => void): IDisposable;
+
+	/**
+	 * Get a contribution of this editor.
+	 * @id Unique identifier of the contribution.
+	 * @return The contribution or null if contribution not found.
+	 */
+	getContribution<T extends INotebookEditorContribution>(id: string): T;
 }
 
 export interface INotebookCellList {

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

@@ -29,9 +29,8 @@ import { BaseEditor } from 'vs/workbench/browser/parts/editor/baseEditor';
 import { IEditorGroupView } from 'vs/workbench/browser/parts/editor/editor';
 import { EditorOptions, IEditorCloseEvent, IEditorMemento } from 'vs/workbench/common/editor';
 import { CELL_MARGIN, CELL_RUN_GUTTER, EDITOR_TOP_MARGIN, EDITOR_TOP_PADDING, EDITOR_BOTTOM_PADDING } from 'vs/workbench/contrib/notebook/browser/constants';
-import { FoldingController } from 'vs/workbench/contrib/notebook/browser/contrib/fold/folding';
 import { NotebookFindWidget } from 'vs/workbench/contrib/notebook/browser/contrib/notebookFindWidget';
-import { CellEditState, CellFocusMode, ICellRange, ICellViewModel, INotebookCellList, INotebookEditor, INotebookEditorMouseEvent, NotebookLayoutInfo, NOTEBOOK_EDITOR_EDITABLE, NOTEBOOK_EDITOR_EXECUTING_NOTEBOOK, NOTEBOOK_EDITOR_FOCUSED } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
+import { CellEditState, CellFocusMode, ICellRange, ICellViewModel, INotebookCellList, INotebookEditor, INotebookEditorMouseEvent, NotebookLayoutInfo, NOTEBOOK_EDITOR_EDITABLE, NOTEBOOK_EDITOR_EXECUTING_NOTEBOOK, NOTEBOOK_EDITOR_FOCUSED, INotebookEditorContribution } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { NotebookEditorInput, NotebookEditorModel } from 'vs/workbench/contrib/notebook/browser/notebookEditorInput';
 import { INotebookService } from 'vs/workbench/contrib/notebook/browser/notebookService';
 import { NotebookCellList } from 'vs/workbench/contrib/notebook/browser/view/notebookCellList';
@@ -45,6 +44,8 @@ import { CellKind, CellUri, IOutput } from 'vs/workbench/contrib/notebook/common
 import { Webview } from 'vs/workbench/contrib/webview/browser/webview';
 import { getExtraColor } from 'vs/workbench/contrib/welcome/walkThrough/common/walkThroughUtils';
 import { IEditorGroup, IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
+import { NotebookEditorExtensionsRegistry } from 'vs/workbench/contrib/notebook/browser/notebookEditorExtensions';
+import { onUnexpectedError } from 'vs/base/common/errors';
 
 const $ = DOM.$;
 const NOTEBOOK_EDITOR_VIEW_STATE_PREFERENCE_KEY = 'NotebookEditorViewState';
@@ -109,7 +110,8 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 	private editorExecutingNotebook: IContextKey<boolean> | null = null;
 	private outputRenderer: OutputRenderer;
 	private findWidget: NotebookFindWidget;
-	private folding: FoldingController | null = null;
+	// private folding: FoldingController | null = null;
+	protected readonly _contributions: { [key: string]: INotebookEditorContribution; };
 
 	constructor(
 		@ITelemetryService telemetryService: ITelemetryService,
@@ -128,6 +130,27 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 		this.outputRenderer = new OutputRenderer(this, this.instantiationService);
 		this.findWidget = this.instantiationService.createInstance(NotebookFindWidget, this);
 		this.findWidget.updateTheme(this.themeService.getColorTheme());
+
+		this._contributions = {};
+		const contributions = NotebookEditorExtensionsRegistry.getEditorContributions();
+
+		for (const desc of contributions) {
+			try {
+				const contribution = this.instantiationService.createInstance(desc.ctor, this);
+				this._contributions[desc.id] = contribution;
+			} catch (err) {
+				onUnexpectedError(err);
+			}
+		}
+	}
+
+	private readonly _onDidChangeModel = new Emitter<void>();
+	readonly onDidChangeModel: Event<void> = this._onDidChangeModel.event;
+
+
+	set viewModel(newModel: NotebookViewModel | undefined) {
+		this.notebookViewModel = newModel;
+		this._onDidChangeModel.fire();
 	}
 
 	get viewModel() {
@@ -365,7 +388,8 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 	private detachModel() {
 		this.localStore.clear();
 		this.list?.detachViewModel();
-		this.notebookViewModel?.dispose();
+		this.viewModel?.dispose();
+		// avoid event
 		this.notebookViewModel = undefined;
 		this.webview?.clearInsets();
 		this.webview?.clearPreloadsCache();
@@ -382,29 +406,35 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 		await this.webview.waitForInitialization();
 
 		this.eventDispatcher = new NotebookEventDispatcher();
-		this.notebookViewModel = this.instantiationService.createInstance(NotebookViewModel, input.viewType!, model, this.eventDispatcher, this.getLayoutInfo());
-		this.editorEditable?.set(!!this.notebookViewModel.metadata?.editable);
+		this.viewModel = this.instantiationService.createInstance(NotebookViewModel, input.viewType!, model, this.eventDispatcher, this.getLayoutInfo());
+		this.editorEditable?.set(!!this.viewModel.metadata?.editable);
 		this.eventDispatcher.emit([new NotebookLayoutChangedEvent({ width: true, fontInfo: true }, this.getLayoutInfo())]);
 
 		this.localStore.add(this.eventDispatcher.onDidChangeMetadata((e) => {
 			this.editorEditable?.set(e.source.editable);
 		}));
 
-		// load contributions
-		this.folding = this.localStore.add(this.instantiationService.createInstance(FoldingController, this));
-
 		// restore view states, including contributions
 		const viewState = this.loadTextEditorViewState(input);
 
 		{
 			// restore view state
-			this.notebookViewModel.restoreEditorViewState(viewState);
+			this.viewModel.restoreEditorViewState(viewState);
 
 			// contribution state restore
-			this.folding?.applyMemento(viewState?.hiddenFoldingRanges || []);
+
+			const contributionsState = viewState?.contributionsState || {};
+			const keys = Object.keys(this._contributions);
+			for (let i = 0, len = keys.length; i < len; i++) {
+				const id = keys[i];
+				const contribution = this._contributions[id];
+				if (typeof contribution.restoreViewState === 'function') {
+					contribution.restoreViewState(contributionsState[id]);
+				}
+			}
 		}
 
-		this.webview?.updateRendererPreloads(this.notebookViewModel.renderers);
+		this.webview?.updateRendererPreloads(this.viewModel.renderers);
 
 		this.localStore.add(this.list!.onWillScroll(e => {
 			this.webview!.updateViewScrollTop(-e.scrollTop, []);
@@ -444,7 +474,7 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 			});
 		}));
 
-		this.list!.attachViewModel(this.notebookViewModel);
+		this.list!.attachViewModel(this.viewModel);
 		this.localStore.add(this.list!.onDidRemoveOutput(output => {
 			this.removeInset(output);
 		}));
@@ -512,11 +542,17 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 			}
 
 			// Save contribution view states
-			if (this.folding) {
-				const foldingState = this.folding.getMemento();
-				state.hiddenFoldingRanges = foldingState;
+			const contributionsState: { [key: string]: any } = {};
+
+			const keys = Object.keys(this._contributions);
+			for (const id of keys) {
+				const contribution = this._contributions[id];
+				if (typeof contribution.saveViewState === 'function') {
+					contributionsState[id] = contribution.saveViewState();
+				}
 			}
 
+			state.contributionsState = contributionsState;
 			this.editorMemento.saveEditorState(this.group, input.resource, state);
 		}
 	}
@@ -906,6 +942,22 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 
 	//#endregion
 
+	//#region Editor Contributions
+	public getContribution<T extends INotebookEditorContribution>(id: string): T {
+		return <T>(this._contributions[id] || null);
+	}
+
+	//#endregion
+	dispose() {
+		const keys = Object.keys(this._contributions);
+		for (let i = 0, len = keys.length; i < len; i++) {
+			const contributionId = keys[i];
+			this._contributions[contributionId].dispose();
+		}
+
+		super.dispose();
+	}
+
 	toJSON(): any {
 		return {
 			notebookHandle: this.viewModel?.handle

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

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { BrandedService, IConstructorSignature1 } from 'vs/platform/instantiation/common/instantiation';
+import { INotebookEditor, INotebookEditorContribution } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
+
+export type INotebookEditorContributionCtor = IConstructorSignature1<INotebookEditor, INotebookEditorContribution>;
+
+
+export interface INotebookEditorContributionDescription {
+	id: string;
+	ctor: INotebookEditorContributionCtor;
+}
+
+class EditorContributionRegistry {
+	public static readonly INSTANCE = new EditorContributionRegistry();
+	private readonly editorContributions: INotebookEditorContributionDescription[];
+
+	constructor() {
+		this.editorContributions = [];
+	}
+
+	public registerEditorContribution<Services extends BrandedService[]>(id: string, ctor: { new(editor: INotebookEditor, ...services: Services): INotebookEditorContribution }): void {
+		this.editorContributions.push({ id, ctor: ctor as INotebookEditorContributionCtor });
+	}
+
+	public getEditorContributions(): INotebookEditorContributionDescription[] {
+		return this.editorContributions.slice(0);
+	}
+}
+
+export function registerNotebookContribution<Services extends BrandedService[]>(id: string, ctor: { new(editor: INotebookEditor, ...services: Services): INotebookEditorContribution }): void {
+	EditorContributionRegistry.INSTANCE.registerEditorContribution(id, ctor);
+}
+
+export namespace NotebookEditorExtensionsRegistry {
+
+	export function getEditorContributions(): INotebookEditorContributionDescription[] {
+		return EditorContributionRegistry.INSTANCE.getEditorContributions();
+	}
+}

src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel.ts

@@ -37,6 +37,7 @@ export interface INotebookEditorViewState {
 	scrollPosition?: { left: number; top: number; };
 	focus?: number;
 	editorFocused?: boolean;
+	contributionsState?: { [id: string]: any };
 }
 
 export interface ICellModelDecorations {

src/vs/workbench/contrib/notebook/test/testNotebookEditor.ts

@@ -9,7 +9,7 @@ import { CellKind, IOutput, CellUri, NotebookCellMetadata } from 'vs/workbench/c
 import { NotebookViewModel, IModelDecorationsChangeAccessor, CellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { NotebookEditorModel } from 'vs/workbench/contrib/notebook/browser/notebookEditorInput';
-import { INotebookEditor, NotebookLayoutInfo, ICellViewModel, ICellRange, INotebookEditorMouseEvent } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
+import { INotebookEditor, NotebookLayoutInfo, ICellViewModel, ICellRange, INotebookEditorMouseEvent, INotebookEditorContribution } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { IMouseWheelEvent } from 'vs/base/browser/mouseEvent';
 import { OutputRenderer } from 'vs/workbench/contrib/notebook/browser/view/output/outputRenderer';
 import { BareFontInfo } from 'vs/editor/common/config/fontInfo';
@@ -21,6 +21,7 @@ import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/mode
 import { NotebookEventDispatcher } from 'vs/workbench/contrib/notebook/browser/viewModel/eventDispatcher';
 import { Webview } from 'vs/workbench/contrib/webview/browser/webview';
 import { IDisposable } from 'vs/base/common/lifecycle';
+import { Emitter, Event } from 'vs/base/common/event';
 
 export class TestCell extends NotebookCellTextModel {
 	constructor(
@@ -50,6 +51,12 @@ export class TestNotebookEditor implements INotebookEditor {
 
 	constructor(
 	) { }
+
+	private _onDidChangeModel = new Emitter<void>();
+	onDidChangeModel: Event<void> = this._onDidChangeModel.event;
+	getContribution<T extends INotebookEditorContribution>(id: string): T {
+		throw new Error('Method not implemented.');
+	}
 	onMouseUp(listener: (e: INotebookEditorMouseEvent) => void): IDisposable {
 		throw new Error('Method not implemented.');
 	}