Improve code selection and undo, add tests (#10160)

src/sql/workbench/contrib/notebook/browser/markdownToolbarActions.ts

@@ -85,30 +85,9 @@ export class MarkdownTextTransformer {
 					let markdownLineType = this.getMarkdownLineType(type);
 					isUndo = this.isUndoOperation(selection, type, markdownLineType, editorModel);
 					if (isUndo) {
-						if (markdownLineType === MarkdownLineType.BEGIN_AND_END_LINES) {
+						this.handleUndoOperation(markdownLineType, startRange, endRange, editorModel, beginInsertedCode, endInsertedCode, selections, selection);
-							startRange = this.getIRangeWithOffsets(startRange, -1 * beginInsertedCode.length, 0, 0, 0);
-							endRange = this.getIRangeWithOffsets(endRange, 0, 0, endInsertedCode.length, 0);
-							editorModel.pushEditOperations(selections, [{ range: endRange, text: '' }, { range: startRange, text: '' }], null);
-						} else {
-							let operations: IIdentifiedSingleEditOperation[] = [];
-							startRange = this.getIRangeWithOffsets(startRange, 0, 0, beginInsertedCode.length, 0);
-							for (let i = 0; i < selection.endLineNumber - selection.startLineNumber + 1; i++) {
-								operations.push({ range: this.transformRangeByLineOffset(startRange, i), text: '' });
-							}
-							editorModel.pushEditOperations(selections, operations, null);
-						}
 					} else {
-						// If the markdown we're inserting only needs to be added to the begin and end lines, add those edit operations directly
+						this.handleTransformOperation(markdownLineType, startRange, endRange, editorModel, beginInsertedCode, endInsertedCode, selections, selection);
-						if (markdownLineType === MarkdownLineType.BEGIN_AND_END_LINES) {
-							editorModel.pushEditOperations(selections, [{ range: startRange, text: beginInsertedCode }, { range: endRange, text: endInsertedCode }], null);
-						} else { // Otherwise, add an operation per line (plus the operation at the last column + line)
-							let operations: IIdentifiedSingleEditOperation[] = [];
-							for (let i = 0; i < selection.endLineNumber - selection.startLineNumber + 1; i++) {
-								operations.push({ range: this.transformRangeByLineOffset(startRange, i), text: beginInsertedCode });
-							}
-							operations.push({ range: endRange, text: endInsertedCode });
-							editorModel.pushEditOperations(selections, operations, null);
-						}
 					}
 				}
 
@@ -116,7 +95,7 @@ export class MarkdownTextTransformer {
 				// Otherwise, the selection will not be correct after the transformation
 				let offset = selection.startLineNumber === selection.endLineNumber ? beginInsertedCode.length : 0;
 				endRange = this.getIRangeWithOffsets(endRange, offset, 0, offset, 0);
-				this.setEndSelection(endRange, type, editorControl, nothingSelected, isUndo);
+				this.setEndSelection(endRange, type, editorControl, editorModel, nothingSelected, isUndo);
 			}
 			// Always give focus back to the editor after pressing the button
 			editorControl.focus();
@@ -182,6 +161,8 @@ export class MarkdownTextTransformer {
 			case MarkdownButtonType.UNORDERED_LIST:
 			case MarkdownButtonType.ORDERED_LIST:
 				return MarkdownLineType.EVERY_LINE;
+			case MarkdownButtonType.CODE:
+				return MarkdownLineType.WRAPPED_ABOVE_AND_BELOW;
 			default:
 				return MarkdownLineType.BEGIN_AND_END_LINES;
 		}
@@ -205,7 +186,7 @@ export class MarkdownTextTransformer {
 
 	private getEditorControl(): CodeEditorWidget | undefined {
 		if (!this._notebookEditor) {
-			this._notebookEditor = this._notebookService.findNotebookEditor(this._cellModel.notebookModel.notebookUri);
+			this._notebookEditor = this._notebookService.findNotebookEditor(this._cellModel?.notebookModel?.notebookUri);
 		}
 		if (this._notebookEditor?.cellEditors?.length > 0) {
 			// Find cell editor provider via cell guid
@@ -232,28 +213,48 @@ export class MarkdownTextTransformer {
 	 * @param editorControl code editor widget
 	 * @param noSelection controls whether there was no previous selection in the editor
 	 */
-	private setEndSelection(endRange: IRange, type: MarkdownButtonType, editorControl: CodeEditorWidget, noSelection: boolean, isUndo: boolean): void {
+	private setEndSelection(endRange: IRange, type: MarkdownButtonType, editorControl: CodeEditorWidget, editorModel: TextModel, noSelection: boolean, isUndo: boolean): void {
 		if (!endRange || !editorControl || isUndo) {
 			return;
 		}
 		let offset = this.getColumnOffsetForSelection(type, noSelection);
 		if (offset > -1) {
-			let newRange: IRange = {
+			let newRange: IRange;
-				startColumn: endRange.startColumn + offset,
+			if (type !== MarkdownButtonType.CODE) {
-				startLineNumber: endRange.startLineNumber,
+				newRange = {
-				endColumn: endRange.startColumn + offset,
+					startColumn: endRange.startColumn + offset,
-				endLineNumber: endRange.endLineNumber
+					startLineNumber: endRange.startLineNumber,
-			};
+					endColumn: endRange.startColumn + offset,
+					endLineNumber: endRange.endLineNumber
+				};
+			} else {
+				newRange = {
+					startColumn: 1,
+					startLineNumber: endRange.startLineNumber + 1,
+					endColumn: 1,
+					endLineNumber: endRange.endLineNumber + 1
+				};
+			}
 			editorControl.setSelection(newRange);
 		} else {
-			if (this.getMarkdownLineType(type) === MarkdownLineType.BEGIN_AND_END_LINES) {
+			let markdownLineType = this.getMarkdownLineType(type);
-				let currentSelection = editorControl.getSelection();
+			let currentSelection = editorControl.getSelection();
+			if (markdownLineType === MarkdownLineType.BEGIN_AND_END_LINES) {
 				editorControl.setSelection({
 					startColumn: currentSelection.startColumn + this.getStartTextToInsert(type).length,
 					startLineNumber: currentSelection.startLineNumber,
 					endColumn: currentSelection.endColumn - this.getEndTextToInsert(type).length,
 					endLineNumber: currentSelection.endLineNumber
 				});
+			} else if (markdownLineType === MarkdownLineType.WRAPPED_ABOVE_AND_BELOW) {
+				// Subtracting 1 because the last line will have the end text (e.g. ``` for code)
+				let endLineLength = editorModel.getLineLength(currentSelection.endLineNumber - 1);
+				editorControl.setSelection({
+					startColumn: 1,
+					startLineNumber: currentSelection.startLineNumber + 1,
+					endColumn: endLineLength + 1,
+					endLineNumber: currentSelection.endLineNumber - 1
+				});
 			}
 		}
 	}
@@ -266,7 +267,7 @@ export class MarkdownTextTransformer {
 	 * @param editorModel text model for the cell
 	 */
 	private isUndoOperation(selection: Selection, type: MarkdownButtonType, lineType: MarkdownLineType, editorModel: TextModel): boolean {
-		if (lineType === MarkdownLineType.BEGIN_AND_END_LINES) {
+		if (lineType === MarkdownLineType.BEGIN_AND_END_LINES || lineType === MarkdownLineType.WRAPPED_ABOVE_AND_BELOW) {
 			let selectedText = this.getExtendedSelectedText(selection, type, lineType, editorModel);
 			return selectedText && selectedText.startsWith(this.getStartTextToInsert(type)) && selectedText.endsWith(this.getEndTextToInsert(type));
 		} else {
@@ -274,6 +275,39 @@ export class MarkdownTextTransformer {
 		}
 	}
 
+	private handleUndoOperation(markdownLineType: MarkdownLineType, startRange: IRange, endRange: IRange, editorModel: TextModel, beginInsertedCode: string, endInsertedCode: string, selections: Selection[], selection: Selection): void {
+		if (markdownLineType === MarkdownLineType.BEGIN_AND_END_LINES) {
+			startRange = this.getIRangeWithOffsets(startRange, -1 * beginInsertedCode.length, 0, 0, 0);
+			endRange = this.getIRangeWithOffsets(endRange, 0, 0, endInsertedCode.length, 0);
+			editorModel.pushEditOperations(selections, [{ range: endRange, text: '' }, { range: startRange, text: '' }], null);
+		} else if (markdownLineType === MarkdownLineType.WRAPPED_ABOVE_AND_BELOW) {
+			startRange = this.getIRangeWithOffsets(startRange, 0, -1, 0, 0);
+			endRange = this.getIRangeWithOffsets(endRange, 0, 0, endInsertedCode.length, 1);
+			editorModel.pushEditOperations(selections, [{ range: endRange, text: '' }, { range: startRange, text: '' }], null);
+		} else {
+			let operations: IIdentifiedSingleEditOperation[] = [];
+			startRange = this.getIRangeWithOffsets(startRange, 0, 0, beginInsertedCode.length, 0);
+			for (let i = 0; i < selection.endLineNumber - selection.startLineNumber + 1; i++) {
+				operations.push({ range: this.transformRangeByLineOffset(startRange, i), text: '' });
+			}
+			editorModel.pushEditOperations(selections, operations, null);
+		}
+	}
+
+	handleTransformOperation(markdownLineType: MarkdownLineType, startRange: IRange, endRange: IRange, editorModel: TextModel, beginInsertedCode: string, endInsertedCode: string, selections: Selection[], selection: Selection): void {
+		// If the markdown we're inserting only needs to be added to the begin and end lines, add those edit operations directly
+		if (markdownLineType === MarkdownLineType.BEGIN_AND_END_LINES || markdownLineType === MarkdownLineType.WRAPPED_ABOVE_AND_BELOW) {
+			editorModel.pushEditOperations(selections, [{ range: startRange, text: beginInsertedCode }, { range: endRange, text: endInsertedCode }], null);
+		} else { // Otherwise, add an operation per line (plus the operation at the last column + line)
+			let operations: IIdentifiedSingleEditOperation[] = [];
+			for (let i = 0; i < selection.endLineNumber - selection.startLineNumber + 1; i++) {
+				operations.push({ range: this.transformRangeByLineOffset(startRange, i), text: beginInsertedCode });
+			}
+			operations.push({ range: endRange, text: endInsertedCode });
+			editorModel.pushEditOperations(selections, operations, null);
+		}
+	}
+
 	/**
 	 * Gets the extended selected text (current selection + potential beginning + ending transformed text)
 	 * @param selection Current selection in editor
@@ -289,6 +323,13 @@ export class MarkdownTextTransformer {
 				endColumn: selection.endColumn + this.getEndTextToInsert(type).length,
 				endLineNumber: selection.endLineNumber
 			});
+		} else if (lineType === MarkdownLineType.WRAPPED_ABOVE_AND_BELOW) {
+			return editorModel.getValueInRange({
+				startColumn: 1,
+				startLineNumber: selection.startLineNumber - 1,
+				endColumn: this.getEndTextToInsert(type).length + 1,
+				endLineNumber: selection.endLineNumber + 1
+			});
 		}
 		return '';
 	}
@@ -300,9 +341,6 @@ export class MarkdownTextTransformer {
 	 * @param editorModel TextModel
 	 */
 	private everyLineMatchesBeginString(selection: Selection, type: MarkdownButtonType, editorModel: TextModel): boolean {
-		if (this.getMarkdownLineType(type) !== MarkdownLineType.EVERY_LINE) {
-			return false;
-		}
 		for (let selectionLine = selection.startLineNumber; selectionLine <= selection.endLineNumber; selectionLine++) {
 			if (!editorModel.getLineContent(selectionLine).startsWith(this.getStartTextToInsert(type))) {
 				return false;
@@ -341,7 +379,9 @@ export enum MarkdownButtonType {
 }
 
 // If ALL_LINES, we need to insert markdown at each line (e.g. lists)
+// WRAPPED_ABOVE_AND_BELOW puts text above and below the highlighted text
 export enum MarkdownLineType {
 	BEGIN_AND_END_LINES,
-	EVERY_LINE
+	EVERY_LINE,
+	WRAPPED_ABOVE_AND_BELOW
 }

src/sql/workbench/contrib/notebook/test/browser/MarkdownTextTransformer.test.ts

@@ -37,9 +37,11 @@ suite('MarkdownTextTransformer', () => {
 	let markdownTextTransformer: MarkdownTextTransformer;
 	let widget: IEditor;
 	let textModel: TextModel;
+	let notebookEditor: TestNotebookEditor;
+	let mockNotebookService: TypeMoq.Mock<INotebookService>;
+	let cellModel: CellModel;
 
 	setup(() => {
-		let mockNotebookService: TypeMoq.Mock<INotebookService>;
 		const dialogService = new TestDialogService();
 		const notificationService = new TestNotificationService();
 		const undoRedoService = new UndoRedoService(dialogService, notificationService);
@@ -55,8 +57,8 @@ suite('MarkdownTextTransformer', () => {
 		mockNotebookService = TypeMoq.Mock.ofType(NotebookService, undefined, new TestLifecycleService(), undefined, undefined, undefined, instantiationService, new MockContextKeyService(),
 			undefined, undefined, undefined, undefined, undefined, undefined, TestEnvironmentService);
 
-		let cellModel = new CellModel(undefined, undefined, mockNotebookService.object);
+		cellModel = new CellModel(undefined, undefined, mockNotebookService.object);
-		let notebookEditor = new TestNotebookEditor(cellModel.cellGuid, instantiationService);
+		notebookEditor = new TestNotebookEditor(cellModel.cellGuid, instantiationService);
 		markdownTextTransformer = new MarkdownTextTransformer(mockNotebookService.object, cellModel, notebookEditor);
 		mockNotebookService.setup(s => s.findNotebookEditor(TypeMoq.It.isAny())).returns(() => notebookEditor);
 
@@ -81,8 +83,8 @@ suite('MarkdownTextTransformer', () => {
 		testWithNoSelection(MarkdownButtonType.BOLD, '');
 		testWithNoSelection(MarkdownButtonType.ITALIC, '__', true);
 		testWithNoSelection(MarkdownButtonType.ITALIC, '');
-		// testWithNoSelection(MarkdownButtonType.CODE, '```\n\n```', true);
+		testWithNoSelection(MarkdownButtonType.CODE, '```\n\n```', true);
-		// testWithNoSelection(MarkdownButtonType.CODE, '\n');
+		testWithNoSelection(MarkdownButtonType.CODE, '');
 		testWithNoSelection(MarkdownButtonType.HIGHLIGHT, '<mark></mark>', true);
 		testWithNoSelection(MarkdownButtonType.HIGHLIGHT, '');
 		testWithNoSelection(MarkdownButtonType.LINK, '[]()', true);
@@ -128,6 +130,17 @@ suite('MarkdownTextTransformer', () => {
 		testWithMultipleLinesSelected(MarkdownButtonType.IMAGE, '![Multi\nLines\nSelected]()');
 	});
 
+	test('Ensure notebook editor returns expected object', () => {
+		assert.deepEqual(notebookEditor, markdownTextTransformer.notebookEditor, 'Notebook editor does not match expected value');
+		// Set markdown text transformer to not have a notebook editor passed in
+		markdownTextTransformer = new MarkdownTextTransformer(mockNotebookService.object, cellModel);
+		assert.equal(markdownTextTransformer.notebookEditor, undefined, 'No notebook editor should be returned');
+		// Even after text is attempted to be transformed, there should be no editor, and therefore nothing on the text model
+		markdownTextTransformer.transformText(MarkdownButtonType.BOLD);
+		assert.equal(markdownTextTransformer.notebookEditor, undefined, 'Notebook model does not have a valid uri, so no editor should be returned');
+		assert.equal(textModel.getValue(), '', 'No text should exist on the textModel');
+	});
+
 	function testWithNoSelection(type: MarkdownButtonType, expectedValue: string, setValue = false): void {
 		if (setValue) {
 			textModel.setValue('');