Notebook Tokenization Fixes (#7375)

* Fix don't like; unclear if grammar necessssary too

* Cleanup and sanity check

* Cleanup and sanity check

* Add test

* Call onBeforeAttached for 3 types of editor models

extensions/notebook/package.json

@@ -161,6 +161,13 @@
         "configuration": "./language-configuration.json"
       }
     ],
+    "grammars": [
+      {
+        "language": "notebook",
+        "scopeName": "source.notebook",
+        "path": "./syntaxes/notebook.tmLanguage.json"
+      }
+    ],
     "menus": {
       "commandPalette": [
         {

extensions/notebook/syntaxes/notebook.tmLanguage.json

@@ -0,0 +1,213 @@
+{
+	"information_for_contributors": [
+		"This file has been converted from https://github.com/Microsoft/vscode-JSON.tmLanguage/blob/master/JSON.tmLanguage",
+		"If you want to provide a fix or improvement, please create a pull request against the original repository.",
+		"Once accepted there, we are happy to receive an update request."
+	],
+	"version": "https://github.com/Microsoft/vscode-JSON.tmLanguage/commit/9bd83f1c252b375e957203f21793316203f61f70",
+	"name": "notebook",
+	"scopeName": "source.notebook",
+	"patterns": [
+		{
+			"include": "#value"
+		}
+	],
+	"repository": {
+		"array": {
+			"begin": "\\[",
+			"beginCaptures": {
+				"0": {
+					"name": "punctuation.definition.array.begin.notebook"
+				}
+			},
+			"end": "\\]",
+			"endCaptures": {
+				"0": {
+					"name": "punctuation.definition.array.end.notebook"
+				}
+			},
+			"name": "meta.structure.array.notebook",
+			"patterns": [
+				{
+					"include": "#value"
+				},
+				{
+					"match": ",",
+					"name": "punctuation.separator.array.notebook"
+				},
+				{
+					"match": "[^\\s\\]]",
+					"name": "invalid.illegal.expected-array-separator.notebook"
+				}
+			]
+		},
+		"comments": {
+			"patterns": [
+				{
+					"begin": "/\\*\\*(?!/)",
+					"captures": {
+						"0": {
+							"name": "punctuation.definition.comment.notebook"
+						}
+					},
+					"end": "\\*/",
+					"name": "comment.block.documentation.notebook"
+				},
+				{
+					"begin": "/\\*",
+					"captures": {
+						"0": {
+							"name": "punctuation.definition.comment.notebook"
+						}
+					},
+					"end": "\\*/",
+					"name": "comment.block.notebook"
+				},
+				{
+					"captures": {
+						"1": {
+							"name": "punctuation.definition.comment.notebook"
+						}
+					},
+					"match": "(//).*$\\n?",
+					"name": "comment.line.double-slash.js"
+				}
+			]
+		},
+		"constant": {
+			"match": "\\b(?:true|false|null)\\b",
+			"name": "constant.language.notebook"
+		},
+		"number": {
+			"match": "(?x)        # turn on extended mode\n  -?        # an optional minus\n  (?:\n    0       # a zero\n    |       # ...or...\n    [1-9]   # a 1-9 character\n    \\d*     # followed by zero or more digits\n  )\n  (?:\n    (?:\n      \\.    # a period\n      \\d+   # followed by one or more digits\n    )?\n    (?:\n      [eE]  # an e character\n      [+-]? # followed by an option +/-\n      \\d+   # followed by one or more digits\n    )?      # make exponent optional\n  )?        # make decimal portion optional",
+			"name": "constant.numeric.notebook"
+		},
+		"object": {
+			"begin": "\\{",
+			"beginCaptures": {
+				"0": {
+					"name": "punctuation.definition.dictionary.begin.notebook"
+				}
+			},
+			"end": "\\}",
+			"endCaptures": {
+				"0": {
+					"name": "punctuation.definition.dictionary.end.notebook"
+				}
+			},
+			"name": "meta.structure.dictionary.notebook",
+			"patterns": [
+				{
+					"comment": "the notebook object key",
+					"include": "#objectkey"
+				},
+				{
+					"include": "#comments"
+				},
+				{
+					"begin": ":",
+					"beginCaptures": {
+						"0": {
+							"name": "punctuation.separator.dictionary.key-value.notebook"
+						}
+					},
+					"end": "(,)|(?=\\})",
+					"endCaptures": {
+						"1": {
+							"name": "punctuation.separator.dictionary.pair.notebook"
+						}
+					},
+					"name": "meta.structure.dictionary.value.notebook",
+					"patterns": [
+						{
+							"comment": "the notebook object value",
+							"include": "#value"
+						},
+						{
+							"match": "[^\\s,]",
+							"name": "invalid.illegal.expected-dictionary-separator.notebook"
+						}
+					]
+				},
+				{
+					"match": "[^\\s\\}]",
+					"name": "invalid.illegal.expected-dictionary-separator.notebook"
+				}
+			]
+		},
+		"string": {
+			"begin": "\"",
+			"beginCaptures": {
+				"0": {
+					"name": "punctuation.definition.string.begin.notebook"
+				}
+			},
+			"end": "\"",
+			"endCaptures": {
+				"0": {
+					"name": "punctuation.definition.string.end.notebook"
+				}
+			},
+			"name": "string.quoted.double.notebook",
+			"patterns": [
+				{
+					"include": "#stringcontent"
+				}
+			]
+		},
+		"objectkey": {
+			"begin": "\"",
+			"beginCaptures": {
+				"0": {
+					"name": "punctuation.support.type.property-name.begin.notebook"
+				}
+			},
+			"end": "\"",
+			"endCaptures": {
+				"0": {
+					"name": "punctuation.support.type.property-name.end.notebook"
+				}
+			},
+			"name": "string.notebook support.type.property-name.notebook",
+			"patterns": [
+				{
+					"include": "#stringcontent"
+				}
+			]
+		},
+		"stringcontent": {
+			"patterns": [
+				{
+					"match": "(?x)                # turn on extended mode\n  \\\\                # a literal backslash\n  (?:               # ...followed by...\n    [\"\\\\/bfnrt]     # one of these characters\n    |               # ...or...\n    u               # a u\n    [0-9a-fA-F]{4}) # and four hex digits",
+					"name": "constant.character.escape.notebook"
+				},
+				{
+					"match": "\\\\.",
+					"name": "invalid.illegal.unrecognized-string-escape.notebook"
+				}
+			]
+		},
+		"value": {
+			"patterns": [
+				{
+					"include": "#constant"
+				},
+				{
+					"include": "#number"
+				},
+				{
+					"include": "#string"
+				},
+				{
+					"include": "#array"
+				},
+				{
+					"include": "#object"
+				},
+				{
+					"include": "#comments"
+				}
+			]
+		}
+	}
+}

src/sql/workbench/parts/notebook/browser/models/notebookInput.ts

@@ -341,10 +341,17 @@ export class NotebookInput extends EditorInput {
 		} else {
 			let textOrUntitledEditorModel: UntitledEditorModel | IEditorModel;
 			if (this.resource.scheme === Schemas.untitled) {
-				textOrUntitledEditorModel = this._untitledEditorModel ? this._untitledEditorModel : await this._textInput.resolve();
+				if (this._untitledEditorModel) {
-			}
+					this._untitledEditorModel.textEditorModel.onBeforeAttached();
-			else {
+					textOrUntitledEditorModel = this._untitledEditorModel;
+				} else {
+					let resolvedInput = await this._textInput.resolve();
+					resolvedInput.textEditorModel.onBeforeAttached();
+					textOrUntitledEditorModel = resolvedInput;
+				}
+			} else {
 				const textEditorModelReference = await this.textModelService.createModelReference(this.resource);
+				textEditorModelReference.object.textEditorModel.onBeforeAttached();
 				textOrUntitledEditorModel = await textEditorModelReference.object.load();
 			}
 			this._model = this.instantiationService.createInstance(NotebookEditorModel, this.resource, textOrUntitledEditorModel);
@@ -385,6 +392,9 @@ export class NotebookInput extends EditorInput {
 	}
 
 	public dispose(): void {
+		if (this._model) {
+			this._model.editorModel.textEditorModel.onBeforeDetached();
+		}
 		this._disposeContainer();
 		super.dispose();
 	}

src/sql/workbench/parts/notebook/browser/models/notebookTextFileModel.ts

@@ -69,17 +69,23 @@ export class NotebookTextFileModel {
 			} else {
 				newOutput = '\n'.concat(newOutput).concat('\n');
 			}
-			let range = this.getEndOfOutputs(textEditorModel, contentChange.cells[0].cellGuid);
+
-			if (range) {
+			// Execution count will always be after the end of the outputs in JSON. This is a sanity mechanism.
+			let executionCountMatch = this.getExecutionCountRange(textEditorModel, contentChange.cells[0].cellGuid);
+			if (!executionCountMatch || !executionCountMatch.range) {
+				return false;
+			}
+
+			let endOutputsRange = this.getEndOfOutputs(textEditorModel, contentChange.cells[0].cellGuid);
+			if (endOutputsRange && endOutputsRange.startLineNumber < executionCountMatch.range.startLineNumber) {
 				textEditorModel.textEditorModel.applyEdits([{
-					range: new Range(range.startLineNumber, range.startColumn, range.startLineNumber, range.startColumn),
+					range: new Range(endOutputsRange.startLineNumber, endOutputsRange.startColumn, endOutputsRange.startLineNumber, endOutputsRange.startColumn),
 					text: newOutput
 				}]);
+				return true;
 			}
-		} else {
-			return false;
 		}
-		return true;
+		return false;
 	}
 
 	public transformAndApplyEditForCellUpdated(contentChange: NotebookContentChange, textEditorModel: TextFileEditorModel | UntitledEditorModel): boolean {

src/sql/workbench/parts/notebook/test/browser/notebookEditorModel.test.ts

@@ -589,6 +589,66 @@ suite('Notebook Editor Model', function (): void {
 
 		should(notebookEditorModel.lastEditFullReplacement).equal(false);
 	});
+	test('should not insert update at incorrect location', async function (): Promise<void> {
+		await createNewNotebookModel();
+		let notebookEditorModel = await createTextEditorModel(this);
+		notebookEditorModel.replaceEntireTextEditorModel(notebookModel, undefined);
+
+		let newCell = notebookModel.addCell(CellTypes.Code);
+
+		let contentChange: NotebookContentChange = {
+			changeType: NotebookChangeType.CellsModified,
+			cells: [newCell],
+			cellIndex: 0
+		};
+		notebookEditorModel.updateModel(contentChange, NotebookChangeType.CellsModified);
+		should(notebookEditorModel.lastEditFullReplacement).equal(true);
+
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(14)).equal('            "outputs": [');
+
+		// First update the model with unmatched brackets
+		let newUnmatchedBracketOutput: nb.IStreamResult = { output_type: 'stream', name: 'stdout', text: '[0em' };
+		newCell[<any>'_outputs'] = newCell.outputs.concat(newUnmatchedBracketOutput);
+
+		contentChange = {
+			changeType: NotebookChangeType.CellOutputUpdated,
+			cells: [newCell]
+		};
+
+		notebookEditorModel.updateModel(contentChange, NotebookChangeType.CellOutputUpdated);
+
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(8)).equal('            "source": [');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(12)).equal('                "azdata_cell_guid": "' + newCell.cellGuid + '"');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(14)).equal('            "outputs": [');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(26)).equal('    "text": "[0em"');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(27)).equal('}');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(28)).equal('            ],');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(29)).equal('            "execution_count": 0');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(30)).equal('        }');
+
+		should(notebookEditorModel.lastEditFullReplacement).equal(false);
+
+		// Now test updating the model after an unmatched bracket was previously output
+		let newBracketlessOutput: nb.IStreamResult = { output_type: 'stream', name: 'stdout', text: 'test test test' };
+		newCell[<any>'_outputs'] = newCell[<any>'_outputs'].concat(newBracketlessOutput);
+
+		contentChange = {
+			changeType: NotebookChangeType.CellOutputUpdated,
+			cells: [newCell]
+		};
+
+		notebookEditorModel.updateModel(contentChange, NotebookChangeType.CellOutputUpdated);
+
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(32)).equal('                    "text": "test test test"');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(33)).equal('                }');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(34)).equal('            ],');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(35)).equal('            "execution_count": 0');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(36)).equal('        }');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(37)).equal('    ]');
+		should(notebookEditorModel.editorModel.textEditorModel.getLineContent(38)).equal('}');
+
+		should(notebookEditorModel.lastEditFullReplacement).equal(true);
+	});
 
 	test('should not replace entire text model for output changes (1st update)', async function (): Promise<void> {
 		await createNewNotebookModel();