Incremental syntax highlight for IRB source lines

Closes: https://github.com/ruby/ruby/pull/2202

2 files changed, 13 insertions, 1 deletions

diff --git a/lib/irb/color.rb b/lib/irb/color.rb
index eb95da8229..30a8fb5973 100644
--- a/lib/irb/color.rb
+++ b/lib/irb/color.rb
@@ -77,13 +77,21 @@ module IRB # :nodoc:
         return code unless colorable?
 
         colored = +''
+        length = 0
         Ripper.lex(code).each do |(_line, _col), token, str, expr|
           if seq = dispatch_seq(token, expr, str)
-            colored << "#{seq.map { |s| "\e[#{s}m" }.join('')}#{str}#{clear}"
+            str.each_line do |line|
+              colored << "#{seq.map { |s| "\e[#{s}m" }.join('')}#{line}#{clear}"
+            end
           else
             colored << str
           end
+          length += str.length
         end
+
+        # give up colorizing incomplete Ripper tokens
+        return code if length != code.length
+
         colored
       end
 
diff --git a/lib/irb/input-method.rb b/lib/irb/input-method.rb
index 412edcce24..bc144dca9c 100644
--- a/lib/irb/input-method.rb
+++ b/lib/irb/input-method.rb
@@ -222,6 +222,10 @@ module IRB
       end
       Reline.completion_append_character = nil
       Reline.completion_proc = IRB::InputCompletor::CompletionProc
+      Reline.output_modifier_proc = proc do |output|
+        next unless IRB::Color.colorable?
+        IRB::Color.colorize_code(output)
+      end
       Reline.dig_perfect_match_proc = IRB::InputCompletor::PerfectMatchedProc
     end