From a88711ff4affb56589b7c4ca35a89eeb505db0ae Mon Sep 17 00:00:00 2001
From: tompng <tomoyapenguin@gmail.com>
Date: Sat, 20 Jul 2024 04:17:45 +0900
Subject: [PATCH] Colorize command input

---
 lib/irb/input-method.rb       |  8 +++++++-
 test/irb/test_input_method.rb | 16 ++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/lib/irb/input-method.rb b/lib/irb/input-method.rb
index f6b8d00e5..82c1e7329 100644
--- 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&.parse_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
index ce317b4b3..078631db8 100644
--- 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 {}