Move grammar override logic into TextEditorRegistry

Signed-off-by: Nathan Sobo <nathan@github.com>

spec/text-editor-registry-spec.js

@@ -85,6 +85,27 @@ describe('TextEditorRegistry', function () {
     })
   })
 
+  describe('.setGrammarOverride and .clearGrammarOverride', function () {
+    it('sets the editor\'s grammar and does not update it based on other criteria until the override is cleared', async function () {
+      registry.maintainGrammar(editor)
+      await atom.packages.activatePackage('language-c')
+      expect(editor.getGrammar().name).toBe('Null Grammar')
+
+      registry.setGrammarOverride(editor, atom.grammars.grammarForScopeName('source.c'))
+      expect(editor.getGrammar().name).toBe('C')
+
+      editor.getBuffer().setPath('file-1.js')
+      await atom.packages.activatePackage('language-javascript')
+      expect(editor.getGrammar().name).toBe('C')
+
+      editor.getBuffer().setPath('file-2.js')
+      expect(editor.getGrammar().name).toBe('C')
+
+      registry.clearGrammarOverride(editor)
+      expect(editor.getGrammar().name).toBe('JavaScript')
+    })
+  })
+
   describe('.maintainConfig(editor)', function () {
     it('does not update editors when config settings change for unrelated scope selectors', async function () {
       await atom.packages.activatePackage('language-javascript')

spec/text-editor-spec.coffee

@@ -5225,24 +5225,8 @@ describe "TextEditor", ->
         editor.setIndentationForBufferRow(0, 2.1)
         expect(editor.getText()).toBe("    1\n\t2")
 
-  describe ".reloadGrammar()", ->
-    beforeEach ->
-      waitsForPromise ->
-        atom.packages.activatePackage('language-coffee-script')
-
-    it "updates the grammar based on grammar overrides", ->
-      expect(editor.getGrammar().name).toBe 'JavaScript'
-      atom.grammars.setGrammarOverrideForPath(editor.getPath(), 'source.coffee')
-      callback = jasmine.createSpy('callback')
-      editor.onDidChangeGrammar(callback)
-      editor.reloadGrammar()
-      expect(editor.getGrammar().name).toBe 'CoffeeScript'
-      expect(callback.callCount).toBe 1
-      expect(callback.argsForCall[0][0]).toBe atom.grammars.grammarForScopeName('source.coffee')
-
   describe "when the editor's grammar has an injection selector", ->
     beforeEach ->
-
       waitsForPromise ->
         atom.packages.activatePackage('language-text')
 

src/text-editor-registry.js

@@ -55,6 +55,7 @@ export default class TextEditorRegistry {
     this.scopesWithConfigSubscriptions = new Set()
     this.editorsWithMaintainedConfig = new Set()
     this.editorsWithMaintainedGrammar = new Set()
+    this.editorGrammarOverrides = new Map()
     this.editorGrammarScores = new WeakMap()
     this.subscriptions.add(
       this.grammarRegistry.onDidAddGrammar(this.grammarAddedOrUpdated),
@@ -106,23 +107,21 @@ export default class TextEditorRegistry {
 
   maintainGrammar (editor) {
     this.editorsWithMaintainedGrammar.add(editor)
+    this.selectGrammarForEditor(editor)
+    this.subscriptions.add(editor.onDidChangePath(() => {
+      this.editorGrammarScores.delete(editor)
+      this.selectGrammarForEditor(editor)
+    }))
+  }
 
-    const assignGrammar = () => {
-      const {grammar, score} = this.grammarRegistry.selectGrammarWithScore(
-        editor.getPath(),
-        editor.getTextInBufferRange(GRAMMAR_SELECTION_RANGE)
-      )
+  setGrammarOverride (editor, grammar) {
+    this.editorGrammarOverrides.set(editor)
+    editor.setGrammar(grammar)
+  }
 
-      if (!grammar) {
-        throw new Error(`No grammar found for path: ${editor.getPath()}`)
-      }
-
-      editor.setGrammar(grammar)
-      this.editorGrammarScores.set(editor, score)
-    }
-
-    assignGrammar()
-    this.subscriptions.add(editor.onDidChangePath(assignGrammar))
+  clearGrammarOverride (editor) {
+    this.editorGrammarOverrides.delete(editor)
+    this.selectGrammarForEditor(editor)
   }
 
   maintainConfig (editor) {
@@ -155,21 +154,44 @@ export default class TextEditorRegistry {
         if (editor.tokenizedBuffer.hasTokenForSelector(grammar.injectionSelector)) {
           editor.tokenizedBuffer.retokenizeLines()
         }
-      } else {
-        const newScore = this.grammarRegistry.getGrammarScore(
+      } else if (!this.editorGrammarOverrides.has(editor)) {
+        const score = this.grammarRegistry.getGrammarScore(
           grammar,
           editor.getPath(),
           editor.getTextInBufferRange(GRAMMAR_SELECTION_RANGE)
         )
+
         let currentScore = this.editorGrammarScores.get(editor)
-        if (currentScore == null || newScore > currentScore) {
-          editor.setGrammar(grammar, newScore)
+        if (currentScore == null || score > currentScore) {
+          editor.setGrammar(grammar, score)
+          this.editorGrammarScores.set(editor, score)
           this.subscribeToSettingsForEditorScope(editor)
         }
       }
     })
   }
 
+  selectGrammarForEditor (editor) {
+    if (this.editorGrammarOverrides.has(editor)) {
+      return
+    }
+
+    const {grammar, score} = this.grammarRegistry.selectGrammarWithScore(
+      editor.getPath(),
+      editor.getTextInBufferRange(GRAMMAR_SELECTION_RANGE)
+    )
+
+    if (!grammar) {
+      throw new Error(`No grammar found for path: ${editor.getPath()}`)
+    }
+
+    const currentScore = this.editorGrammarScores.get(editor)
+    if (currentScore == null || score > currentScore) {
+      editor.setGrammar(grammar)
+      this.editorGrammarScores.set(editor, score)
+    }
+  }
+
   subscribeToSettingsForEditorScope (editor) {
     const scopeDescriptor = editor.getRootScopeDescriptor()
     const scopeChain = scopeDescriptor.getScopeChain()