From eccb77aca330a732b95f47110574bfc28026f9b0 Mon Sep 17 00:00:00 2001 From: Cory Rivera Date: Sun, 6 Mar 2022 21:35:57 -0800 Subject: [PATCH] Preserve previous code cell's language when creating a new code cell from an existing context. (#18646) --- .../notebook/browser/cellToolbarActions.ts | 4 +-- .../notebook/browser/notebookActions.ts | 2 +- .../electron-browser/notebookModel.test.ts | 27 +++++++++++++++++++ .../workbench/contrib/notebook/test/stubs.ts | 2 +- .../browser/models/modelInterfaces.ts | 2 +- .../notebook/browser/models/notebookModel.ts | 15 ++++++----- 6 files changed, 41 insertions(+), 11 deletions(-) diff --git a/src/sql/workbench/contrib/notebook/browser/cellToolbarActions.ts b/src/sql/workbench/contrib/notebook/browser/cellToolbarActions.ts index 3156cd1f23..c8aeffcbd7 100644 --- a/src/sql/workbench/contrib/notebook/browser/cellToolbarActions.ts +++ b/src/sql/workbench/contrib/notebook/browser/cellToolbarActions.ts @@ -75,7 +75,7 @@ export class SplitCellAction extends CellActionBase { doRun(context: CellContext): Promise { let model = context.model; let index = model.cells.findIndex((cell) => cell.id === context.cell.id); - context.model?.splitCell(context.cell.cellType, this.notebookService, index); + context.model?.splitCell(context.cell.cellType, this.notebookService, index, context.cell.metadata?.language); return Promise.resolve(); } public setListener(context: CellContext) { @@ -246,7 +246,7 @@ export class AddCellFromContextAction extends CellActionBase { if (index !== undefined && this.isAfter) { index += 1; } - model.addCell(this.cellType, index); + model.addCell(this.cellType, index, context.cell.metadata?.language); } catch (error) { let message = getErrorMessage(error); diff --git a/src/sql/workbench/contrib/notebook/browser/notebookActions.ts b/src/sql/workbench/contrib/notebook/browser/notebookActions.ts index 2c69abb7fa..c85d0a37ca 100644 --- a/src/sql/workbench/contrib/notebook/browser/notebookActions.ts +++ b/src/sql/workbench/contrib/notebook/browser/notebookActions.ts @@ -74,7 +74,7 @@ export class AddCellAction extends Action { } } if (context?.model) { - context.model.addCell(this.cellType, index); + context.model.addCell(this.cellType, index, context.cell.metadata?.language); context.model.sendNotebookTelemetryActionEvent(TelemetryKeys.NbTelemetryAction.AddCell, { cell_type: this.cellType }); } } else { diff --git a/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts b/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts index 66110cb387..8f4987e9cd 100644 --- a/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts +++ b/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts @@ -1006,6 +1006,33 @@ suite('notebook model', function (): void { assert.strictEqual(model.languageInfo.name, 'fake', 'Notebook language info is not set properly'); }); + test('Should add new cell with provided cell language', async function () { + let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentLoader); + mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent)); + defaultModelOptions.contentLoader = mockContentManager.object; + + let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService, undoRedoService); + await model.loadContents(); + + const newLanguage = 'CustomCellLanguage'; + let cell = model.addCell(CellTypes.Code, undefined, newLanguage); + assert.strictEqual(model.cells.length, expectedNotebookContent.cells.length + 1, 'New cell was not added to list of cells.'); + assert.strictEqual(cell.language, newLanguage); + }); + + test('Should add new cell with default language if none is provided', async function () { + let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentLoader); + mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent)); + defaultModelOptions.contentLoader = mockContentManager.object; + + let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService, undoRedoService); + await model.loadContents(); + + let cell = model.addCell(CellTypes.Code); + assert.strictEqual(model.cells.length, expectedNotebookContent.cells.length + 1, 'New cell was not added to list of cells.'); + assert.strictEqual(cell.language, expectedNotebookContent.metadata.language_info.name); + }); + async function loadModelAndStartClientSession(notebookContent: nb.INotebookContents): Promise { let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentLoader); mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(notebookContent)); diff --git a/src/sql/workbench/contrib/notebook/test/stubs.ts b/src/sql/workbench/contrib/notebook/test/stubs.ts index 5626398abf..48aec72f13 100644 --- a/src/sql/workbench/contrib/notebook/test/stubs.ts +++ b/src/sql/workbench/contrib/notebook/test/stubs.ts @@ -122,7 +122,7 @@ export class NotebookModelStub implements INotebookModel { getMetaValue(key: string) { throw new Error('Method not implemented.'); } - addCell(cellType: CellType, index?: number): void { + addCell(cellType: CellType, index?: number, language?: string): void { throw new Error('Method not implemented.'); } moveCell(cellModel: ICellModel, direction: MoveDirection): void { diff --git a/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts b/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts index 59fffe5640..951a0f8dbf 100644 --- a/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts +++ b/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts @@ -389,7 +389,7 @@ export interface INotebookModel { /** * Adds a cell to the index of the model */ - addCell(cellType: CellType, index?: number): void; + addCell(cellType: CellType, index?: number, language?: string): void; /** * Moves a cell up/down diff --git a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts index a12a2990b5..efa48cfd79 100644 --- a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts +++ b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts @@ -543,15 +543,15 @@ export class NotebookModel extends Disposable implements INotebookModel { return this._cells.findIndex(cell => cell.equals(cellModel)); } - public addCell(cellType: CellType, index?: number): ICellModel | undefined { + public addCell(cellType: CellType, index?: number, language?: string): ICellModel | undefined { if (this.inErrorState) { return undefined; } - let cell = this.createCell(cellType); + let cell = this.createCell(cellType, language); return this.insertCell(cell, index, true); } - public splitCell(cellType: CellType, notebookService: INotebookService, index?: number, addToUndoStack: boolean = true): ICellModel | undefined { + public splitCell(cellType: CellType, notebookService: INotebookService, index?: number, language?: string, addToUndoStack: boolean = true): ICellModel | undefined { if (this.inErrorState) { return undefined; } @@ -626,7 +626,7 @@ export class NotebookModel extends Disposable implements INotebookModel { } //If the selection is not from the start of the cell, create a new cell. if (headContent.length) { - newCell = this.createCell(cellType); + newCell = this.createCell(cellType, language); newCell.source = newSource; newCellIndex++; this.insertCell(newCell, newCellIndex, false); @@ -639,7 +639,7 @@ export class NotebookModel extends Disposable implements INotebookModel { if (tailCellContent.length) { //tail cell will be of original cell type. - tailCell = this.createCell(this._cells[index].cellType); + tailCell = this.createCell(this._cells[index].cellType, language); let tailSource = source.slice(tailRange.startLineNumber - 1) as string[]; if (selection.endColumn > 1) { partialSource = source.slice(tailRange.startLineNumber - 1, tailRange.startLineNumber)[0].slice(tailRange.startColumn - 1); @@ -833,13 +833,16 @@ export class NotebookModel extends Disposable implements INotebookModel { } } - private createCell(cellType: CellType): ICellModel { + private createCell(cellType: CellType, language?: string): ICellModel { let singleCell: nb.ICellContents = { cell_type: cellType, source: '', metadata: {}, execution_count: undefined }; + if (language) { + singleCell.metadata.language = language; + } return this._notebookOptions.factory.createCell(singleCell, { notebook: this, isTrusted: true }); }