Colorize command input

ruby · Jul 20, 2024 · e62cfc7 · e62cfc7
1 parent 632da0f
commit e62cfc7
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/lib/irb/input-method.rb b/lib/irb/input-method.rb
@@ -270,7 +270,13 @@ def initialize(completor)
           proc do |output, complete: |
             next unless IRB::Color.colorable?
             lvars = IRB.CurrentContext&.local_variables || []
-            IRB::Color.colorize_code(output, complete: complete, local_variables: lvars)
+            if IRB.CurrentContext&.irb&.command?(output)
+              name, sep, arg = output.split(/(\s+)/, 2)
+              arg = IRB::Color.colorize_code(arg, complete: complete, local_variables: lvars)
+              "#{IRB::Color.colorize(name, [:BOLD])}\e[m#{sep}#{arg}"
+            else
+              IRB::Color.colorize_code(output, complete: complete, local_variables: lvars)
+            end
           end
         else
           proc do |output|

diff --git a/test/irb/test_input_method.rb b/test/irb/test_input_method.rb
@@ -8,6 +8,7 @@ module TestIRB
   class InputMethodTest < TestCase
     def setup
       @conf_backup = IRB.conf.dup
+      IRB.init_config(nil)
       IRB.conf[:LC_MESSAGES] = IRB::Locale.new
       save_encodings
     end
@@ -33,6 +34,21 @@ def test_initialization
       assert_not_nil Reline.dig_perfect_match_proc
     end
 
+    def test_colorize
+      original_colorable = IRB::Color.method(:colorable?)
+      IRB::Color.instance_eval { undef :colorable? }
+      IRB::Color.define_singleton_method(:colorable?) { true }
+      workspace = IRB::WorkSpace.new(binding)
+      input_method = IRB::RelineInputMethod.new(IRB::RegexpCompletor.new)
+      IRB.conf[:MAIN_CONTEXT] = IRB::Irb.new(workspace, input_method).context
+      assert_equal "\e[1m$\e[0m\e[m", Reline.output_modifier_proc.call('$', complete: false)
+      assert_equal "\e[1m$\e[0m\e[m  \e[34m\e[1m1\e[0m + \e[34m\e[1m2\e[0m", Reline.output_modifier_proc.call('$  1 + 2', complete: false)
+      assert_equal "\e[32m\e[1m$a\e[0m", Reline.output_modifier_proc.call('$a', complete: false)
+    ensure
+      IRB::Color.instance_eval { undef :colorable? }
+      IRB::Color.define_singleton_method(:colorable?, original_colorable)
+    end
+
     def test_initialization_without_use_autocomplete
       original_show_doc_proc = Reline.dialog_proc(:show_doc)&.dialog_proc
       empty_proc = Proc.new {}