1 files changed, 22 insertions, 5 deletions
diff --git a/lib/gitlab/highlight.rb b/lib/gitlab/highlight.rb
index 280120b0f9e..41296415e35 100644
--- a/lib/gitlab/highlight.rb
+++ b/lib/gitlab/highlight.rb
@@ -1,7 +1,7 @@
 module Gitlab
   class Highlight
-    def self.highlight(blob_name, blob_content, nowrap: true, plain: false)
-      new(blob_name, blob_content, nowrap: nowrap).
+    def self.highlight(blob_name, blob_content, repository: nil, nowrap: true, plain: false)
+      new(blob_name, blob_content, nowrap: nowrap, repository: repository).
         highlight(blob_content, continue: false, plain: plain)
     end
 
@@ -10,12 +10,21 @@ module Gitlab
       return [] unless blob
 
       blob.load_all_data!(repository)
-      highlight(file_name, blob.data).lines.map!(&:html_safe)
+      highlight(file_name, blob.data, repository: repository).lines.map!(&:html_safe)
     end
 
-    def initialize(blob_name, blob_content, nowrap: true)
+    attr_reader :lexer
+    def initialize(blob_name, blob_content, repository: nil, nowrap: true)
+      @blob_name = blob_name
+      @blob_content = blob_content
+      @repository = repository
       @formatter = rouge_formatter(nowrap: nowrap)
-      @lexer = Rouge::Lexer.guess(filename: blob_name, source: blob_content).new rescue Rouge::Lexers::PlainText
+
+      @lexer = custom_language || begin
+        Rouge::Lexer.guess(filename: blob_name, source: blob_content).new
+      rescue Rouge::Lexer::AmbiguousGuess => e
+        e.alternatives.sort_by(&:tag).first
+      end
     end
 
     def highlight(text, continue: true, plain: false)
@@ -30,6 +39,14 @@ module Gitlab
 
     private
 
+    def custom_language
+      language_name = @repository && @repository.gitattribute(@blob_name, 'gitlab-language')
+
+      return nil unless language_name
+
+      Rouge::Lexer.find_fancy(language_name)
+    end
+
     def rouge_formatter(options = {})
       options = options.reverse_merge(
         nowrap: true,