Refactor Gitlab::Highlight and fix tests

1 files changed, 20 insertions, 6 deletions

diff --git a/lib/gitlab/highlight.rb b/lib/gitlab/highlight.rb
index a5b041687e3..28cfebef968 100644
--- a/lib/gitlab/highlight.rb
+++ b/lib/gitlab/highlight.rb
@@ -1,10 +1,7 @@
 module Gitlab
   class Highlight
-    def self.highlight(blob_name, blob_content, nowrap: true, continue: false)
-      formatter = rouge_formatter(nowrap: nowrap)
-
-      lexer = Rouge::Lexer.guess(filename: blob_name, source: blob_content).new rescue Rouge::Lexers::PlainText
-      formatter.format(lexer.lex(blob_content, continue: continue)).html_safe
+    def self.highlight(blob_name, blob_content, nowrap: true)
+      new(blob_name, blob_content, nowrap: nowrap).highlight(blob_content, continue: false)
     end
 
     def self.highlight_lines(repository, ref, file_name)
@@ -14,9 +11,26 @@ module Gitlab
       highlight(file_name, blob.data).lines.map!(&:html_safe)
     end
 
+    def initialize(blob_name, blob_content, nowrap: true)
+      @formatter = rouge_formatter(nowrap: nowrap)
+      @lexer = Rouge::Lexer.guess(filename: blob_name, source: blob_content).new rescue Rouge::Lexers::PlainText
+    end
+
+    def highlight(text, continue: true)
+      @formatter.format(lex(text, continue: continue)).html_safe
+    end
+
     private
 
-    def self.rouge_formatter(options = {})
+    def lex(text, continue: true)
+      if @lexer == Rouge::Lexers::PlainText
+        @lexer.lex(text)
+      else
+        @lexer.lex(text, continue: continue)
+      end
+    end
+
+    def rouge_formatter(options = {})
       options = options.reverse_merge(
         nowrap: true,
         cssclass: 'code highlight',