Colorize IRB's code_around_binding

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

2 files changed, 41 insertions, 2 deletions

diff --git a/test/irb/test_color.rb b/test/irb/test_color.rb
new file mode 100644
index 0000000000..0c41613f51
--- /dev/null
+++ b/test/irb/test_color.rb
@@ -0,0 +1,30 @@
+# frozen_string_literal: false
+require 'test/unit'
+require 'irb/color'
+
+module TestIRB
+  class TestColor < Test::Unit::TestCase
+    CLEAR     = "\e[0m"
+    BOLD      = "\e[1m"
+    UNDERLINE = "\e[4m"
+    RED       = "\e[31m"
+    GREEN     = "\e[32m"
+    BLUE      = "\e[34m"
+    MAGENTA   = "\e[35m"
+    CYAN      = "\e[36m"
+
+    def test_colorize_code
+      {
+        "1" => "#{BLUE}#{BOLD}1#{CLEAR}",
+        "2.3" => "#{MAGENTA}#{BOLD}2.3#{CLEAR}",
+        "['foo', :bar]" => "[#{RED}'#{CLEAR}#{RED}foo#{CLEAR}#{RED}'#{CLEAR}, #{BLUE}#{BOLD}:#{CLEAR}#{BLUE}#{BOLD}bar#{CLEAR}]",
+        "class A; end" => "#{GREEN}class#{CLEAR} #{BLUE}#{BOLD}#{UNDERLINE}A#{CLEAR}; #{GREEN}end#{CLEAR}",
+        "def self.foo; bar; end" => "#{GREEN}def#{CLEAR} #{CYAN}#{BOLD}self#{CLEAR}.#{BLUE}#{BOLD}foo#{CLEAR}; bar; #{GREEN}end#{CLEAR}",
+        'ERB.new("a#{nil}b", trim_mode: "-")' => "#{BLUE}#{BOLD}#{UNDERLINE}ERB#{CLEAR}.new(#{RED}\"#{CLEAR}#{RED}a#{CLEAR}#{RED}\#{#{CLEAR}#{CYAN}#{BOLD}nil#{CLEAR}#{RED}}#{CLEAR}#{RED}b#{CLEAR}#{RED}\"#{CLEAR}, #{MAGENTA}trim_mode:#{CLEAR} #{RED}\"#{CLEAR}#{RED}-#{CLEAR}#{RED}\"#{CLEAR})",
+        "# comment" => "#{BLUE}#{BOLD}# comment#{CLEAR}",
+      }.each do |code, result|
+        assert_equal(result, IRB::Color.colorize_code(code))
+      end
+    end
+  end
+end
diff --git a/test/irb/test_workspace.rb b/test/irb/test_workspace.rb
index 0795b17e09..9c87468cf7 100644
--- a/test/irb/test_workspace.rb
+++ b/test/irb/test_workspace.rb
@@ -2,6 +2,7 @@
 require 'test/unit'
 require 'tempfile'
 require 'irb/workspace'
+require 'irb/color'
 
 module TestIRB
   class TestWorkSpace < Test::Unit::TestCase
@@ -18,7 +19,7 @@ module TestIRB
         f.close
 
         workspace = eval(code, binding, f.path)
-        assert_equal(<<~EOS, workspace.code_around_binding)
+        assert_equal(<<~EOS, without_term { workspace.code_around_binding })
 
           From: #{f.path} @ line 3 :
 
@@ -55,7 +56,7 @@ module TestIRB
           script_lines[f.path] = code.split(/^/)
 
           workspace = eval(code, binding, f.path)
-          assert_equal(<<~EOS, workspace.code_around_binding)
+          assert_equal(<<~EOS, without_term { workspace.code_around_binding })
 
             From: #{f.path} @ line 1 :
 
@@ -90,5 +91,13 @@ module TestIRB
         const_set(:SCRIPT_LINES__, script_lines) if script_lines
       end
     end
+
+    def without_term
+      env = ENV.to_h.dup
+      ENV.delete('TERM')
+      yield
+    ensure
+      ENV.replace(env)
+    end
   end
 end