Highlight the method name in method calls

lib/irb/color.rb

@@ -102,16 +102,17 @@ module Color
       STRING_END:         [RED, BOLD],
       __END__:            [GREEN],
       # tokens from syntax tree traversal
-      method_name:        [BLUE, BOLD],
+      method_name:        [CYAN, BOLD],
+      message_name:       [CYAN],
       symbol:             [YELLOW],
       # special colorization
       error:              [RED, REVERSE],
-      const_env:          [CYAN, BOLD],
     }.transform_values do |styles|
       styles.map { |style| "\e[#{style}m" }.join
     end
     CLEAR_SEQ = "\e[#{CLEAR}m"
-    private_constant :TOKEN_SEQS, :CLEAR_SEQ
+    OPERATORS = %i(!= !~ =~ == === <=> > >= < <= & | ^ >> << - + % / * ** -@ +@ ~ ! [] []=)
+    private_constant :TOKEN_SEQS, :CLEAR_SEQ, :OPERATORS
 
     class << self
       def colorable?
@@ -160,7 +161,7 @@ def colorize(text, seq, colorable: colorable?)
       # If `complete` is false (code is incomplete), this does not warn compile_error.
       # This option is needed to avoid warning a user when the compile_error is happening
       # because the input is not wrong but just incomplete.
-      def colorize_code(code, complete: true, ignore_error: false, colorable: colorable?, local_variables: [])
+      def colorize_code(code, complete: true, ignore_error: false, colorable: colorable?, colorize_call: true, local_variables: [])
         return code unless colorable
 
         result = Prism.parse_lex(code, scopes: [local_variables])
@@ -175,7 +176,7 @@ def colorize_code(code, complete: true, ignore_error: false, colorable: colorabl
           errors = filter_incomplete_code_errors(errors, prism_tokens)
         end
 
-        visitor = ColorizeVisitor.new
+        visitor = ColorizeVisitor.new(colorize_call: colorize_call)
         prism_node.accept(visitor)
 
         error_tokens = errors.map { |e| [e.location.start_line, e.location.start_column, 0, e.location.end_line, e.location.end_column, :error, e.location.slice] }
@@ -203,9 +204,7 @@ def colorize_code(code, complete: true, ignore_error: false, colorable: colorabl
           next if start_line - 1 < line_index || (start_line - 1 == line_index && start_column < col)
 
           flush.call(start_line - 1, start_column)
-          if type == :CONSTANT && value == 'ENV'
-            color = TOKEN_SEQS[:const_env]
-          elsif type == :__END__
+          if type == :__END__
             color = TOKEN_SEQS[type]
             end_line = start_line
             value = '__END__'
@@ -229,7 +228,8 @@ def colorize_code(code, complete: true, ignore_error: false, colorable: colorabl
 
       class ColorizeVisitor < Prism::Visitor
         attr_reader :tokens
-        def initialize
+        def initialize(colorize_call: true)
+          @colorize_call = colorize_call
           @tokens = []
         end
 
@@ -252,6 +252,33 @@ def visit_def_node(node)
           super
         end
 
+        def visit_alias_method_node(node)
+          dispatch_alias_method_name node.new_name
+          dispatch_alias_method_name node.old_name
+          super
+        end
+
+        def visit_call_node(node)
+          if @colorize_call
+            if node.call_operator_loc.nil? && OPERATORS.include?(node.name)
+              # Operators should not be colored as method call
+            elsif (node.call_operator_loc.nil? || node.call_operator_loc.slice == "::") &&
+                /\A\p{Upper}/.match?(node.name)
+              # Constant-like methods should not be colored as method call
+            else
+              dispatch node.message_loc, :message_name
+            end
+          end
+          super
+        end
+
+        def visit_call_operator_write_node(node)
+          dispatch node.message_loc, :message_name if @colorize_call
+          super
+        end
+        alias visit_call_and_write_node visit_call_operator_write_node
+        alias visit_call_or_write_node visit_call_operator_write_node
+
         def visit_interpolated_symbol_node(node)
           dispatch node.opening_loc, :symbol
           node.parts.each do |part|
@@ -279,6 +306,14 @@ def visit_symbol_node(node)
             dispatch node.closing_loc, :symbol
           end
         end
+
+        private
+
+        def dispatch_alias_method_name(node)
+          if node.type == :symbol_node && node.opening_loc.nil?
+            dispatch node.value_loc, :method_name
+          end
+        end
       end
 
       private

lib/irb/color_printer.rb

@@ -49,7 +49,7 @@ def text(str, width = nil)
       when /\A#</, '=', '>'
         super(@colorize ? Color.colorize(str, [:GREEN]) : str, width)
       else
-        super(@colorize ? Color.colorize_code(str, ignore_error: true) : str, width)
+        super(@colorize ? Color.colorize_code(str, ignore_error: true, colorize_call: false) : str, width)
       end
     end
   end

test/irb/helper.rb

@@ -174,7 +174,7 @@ def run_ruby_file(timeout: TIMEOUT_SEC, via_irb: false, &block)
             lines << line
 
             # means the breakpoint is triggered
-            if line.match?(/binding\.irb/)
+            if line.match?(/binding(?<color>\e\[\d+m)?\.\g<color>?irb/)
               while command = @commands.shift
                 write.puts(command)
               end

test/irb/test_color.rb

@@ -61,13 +61,15 @@ def test_colorize_code
         "8i" => "#{BLUE}#{BOLD}8i#{CLEAR}",
         "['foo', :bar]" => "[#{RED}#{BOLD}'#{CLEAR}#{RED}foo#{CLEAR}#{RED}#{BOLD}'#{CLEAR}, #{YELLOW}:#{CLEAR}#{YELLOW}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 = ERB.new("a#{nil}b", trim_mode: "-")' => "erb = #{BLUE}#{BOLD}#{UNDERLINE}ERB#{CLEAR}.new(#{RED}#{BOLD}\"#{CLEAR}#{RED}a#{CLEAR}#{RED}\#{#{CLEAR}#{CYAN}#{BOLD}nil#{CLEAR}#{RED}}#{CLEAR}#{RED}b#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}, #{MAGENTA}trim_mode:#{CLEAR} #{RED}#{BOLD}\"#{CLEAR}#{RED}-#{CLEAR}#{RED}#{BOLD}\"#{CLEAR})",
+        "def self.foo; bar; end" => "#{GREEN}def#{CLEAR} #{CYAN}#{BOLD}self#{CLEAR}.#{CYAN}#{BOLD}foo#{CLEAR}; #{CYAN}bar#{CLEAR}; #{GREEN}end#{CLEAR}",
+        'erb = ERB.new("a#{nil}b", trim_mode: "-")' => "erb = #{BLUE}#{BOLD}#{UNDERLINE}ERB#{CLEAR}.#{CYAN}new#{CLEAR}(#{RED}#{BOLD}\"#{CLEAR}#{RED}a#{CLEAR}#{RED}\#{#{CLEAR}#{CYAN}#{BOLD}nil#{CLEAR}#{RED}}#{CLEAR}#{RED}b#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}, #{MAGENTA}trim_mode:#{CLEAR} #{RED}#{BOLD}\"#{CLEAR}#{RED}-#{CLEAR}#{RED}#{BOLD}\"#{CLEAR})",
         "# comment" => "#{BLUE}#{BOLD}# comment#{CLEAR}",
-        "def f;yield(hello);end" => "#{GREEN}def#{CLEAR} #{BLUE}#{BOLD}f#{CLEAR};#{GREEN}yield#{CLEAR}(hello);#{GREEN}end#{CLEAR}",
+        "def f;yield(hello);end" => "#{GREEN}def#{CLEAR} #{CYAN}#{BOLD}f#{CLEAR};#{GREEN}yield#{CLEAR}(#{CYAN}hello#{CLEAR});#{GREEN}end#{CLEAR}",
+        "alias foo bar" => "#{GREEN}alias#{CLEAR} #{CYAN}#{BOLD}foo#{CLEAR} #{CYAN}#{BOLD}bar#{CLEAR}",
+        "alias :foo :bar" => "#{GREEN}alias#{CLEAR} #{YELLOW}:#{CLEAR}#{YELLOW}foo#{CLEAR} #{YELLOW}:#{CLEAR}#{YELLOW}bar#{CLEAR}",
         '"##@var]"' => "#{RED}#{BOLD}\"#{CLEAR}#{RED}\##{CLEAR}#{RED}\##{CLEAR}@var#{RED}]#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}",
-        '"foo#{a} #{b}"' => "#{RED}#{BOLD}\"#{CLEAR}#{RED}foo#{CLEAR}#{RED}\#{#{CLEAR}a#{RED}}#{CLEAR}#{RED} #{CLEAR}#{RED}\#{#{CLEAR}b#{RED}}#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}",
-        '/r#{e}g/' => "#{RED}#{BOLD}/#{CLEAR}#{RED}r#{CLEAR}#{RED}\#{#{CLEAR}e#{RED}}#{CLEAR}#{RED}g#{CLEAR}#{RED}#{BOLD}/#{CLEAR}",
+        '"foo#{a} #{b}"' => "#{RED}#{BOLD}\"#{CLEAR}#{RED}foo#{CLEAR}#{RED}\#{#{CLEAR}#{CYAN}a#{CLEAR}#{RED}}#{CLEAR}#{RED} #{CLEAR}#{RED}\#{#{CLEAR}#{CYAN}b#{CLEAR}#{RED}}#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}",
+        '/r#{e}g/' => "#{RED}#{BOLD}/#{CLEAR}#{RED}r#{CLEAR}#{RED}\#{#{CLEAR}#{CYAN}e#{CLEAR}#{RED}}#{CLEAR}#{RED}g#{CLEAR}#{RED}#{BOLD}/#{CLEAR}",
         "'a\nb'" => "#{RED}#{BOLD}'#{CLEAR}#{RED}a#{CLEAR}\n#{RED}b#{CLEAR}#{RED}#{BOLD}'#{CLEAR}",
         "%[str]" => "#{RED}#{BOLD}%[#{CLEAR}#{RED}str#{CLEAR}#{RED}#{BOLD}]#{CLEAR}",
         "%Q[str]" => "#{RED}#{BOLD}%Q[#{CLEAR}#{RED}str#{CLEAR}#{RED}#{BOLD}]#{CLEAR}",
@@ -83,40 +85,56 @@ def test_colorize_code
         ":Struct" => "#{YELLOW}:#{CLEAR}#{YELLOW}Struct#{CLEAR}",
         '"#{}"' => "#{RED}#{BOLD}\"#{CLEAR}#{RED}\#{#{CLEAR}#{RED}}#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}",
         ':"a#{}b"' => "#{YELLOW}:\"#{CLEAR}#{YELLOW}a#{CLEAR}#{YELLOW}\#{#{CLEAR}#{YELLOW}}#{CLEAR}#{YELLOW}b#{CLEAR}#{YELLOW}\"#{CLEAR}",
-        ':"a#{ def b; end; \'c\' + "#{ :d }" }e"' => "#{YELLOW}:\"#{CLEAR}#{YELLOW}a#{CLEAR}#{YELLOW}\#{#{CLEAR} #{GREEN}def#{CLEAR} #{BLUE}#{BOLD}b#{CLEAR}; #{GREEN}end#{CLEAR}; #{RED}#{BOLD}'#{CLEAR}#{RED}c#{CLEAR}#{RED}#{BOLD}'#{CLEAR} + #{RED}#{BOLD}\"#{CLEAR}#{RED}\#{#{CLEAR} #{YELLOW}:#{CLEAR}#{YELLOW}d#{CLEAR} #{RED}}#{CLEAR}#{RED}#{BOLD}\"#{CLEAR} #{YELLOW}}#{CLEAR}#{YELLOW}e#{CLEAR}#{YELLOW}\"#{CLEAR}",
+        ':"a#{ def b; end; \'c\' + "#{ :d }" }e"' => "#{YELLOW}:\"#{CLEAR}#{YELLOW}a#{CLEAR}#{YELLOW}\#{#{CLEAR} #{GREEN}def#{CLEAR} #{CYAN}#{BOLD}b#{CLEAR}; #{GREEN}end#{CLEAR}; #{RED}#{BOLD}'#{CLEAR}#{RED}c#{CLEAR}#{RED}#{BOLD}'#{CLEAR} + #{RED}#{BOLD}\"#{CLEAR}#{RED}\#{#{CLEAR} #{YELLOW}:#{CLEAR}#{YELLOW}d#{CLEAR} #{RED}}#{CLEAR}#{RED}#{BOLD}\"#{CLEAR} #{YELLOW}}#{CLEAR}#{YELLOW}e#{CLEAR}#{YELLOW}\"#{CLEAR}",
         "[__FILE__, __LINE__, __ENCODING__]" => "[#{CYAN}#{BOLD}__FILE__#{CLEAR}, #{CYAN}#{BOLD}__LINE__#{CLEAR}, #{CYAN}#{BOLD}__ENCODING__#{CLEAR}]",
         ":self" => "#{YELLOW}:#{CLEAR}#{YELLOW}self#{CLEAR}",
         ":class" => "#{YELLOW}:#{CLEAR}#{YELLOW}class#{CLEAR}",
         "[:end, 2]" => "[#{YELLOW}:#{CLEAR}#{YELLOW}end#{CLEAR}, #{BLUE}#{BOLD}2#{CLEAR}]",
         "[:>, 3]" => "[#{YELLOW}:#{CLEAR}#{YELLOW}>#{CLEAR}, #{BLUE}#{BOLD}3#{CLEAR}]",
         "[:`, 4]" => "[#{YELLOW}:#{CLEAR}#{YELLOW}`#{CLEAR}, #{BLUE}#{BOLD}4#{CLEAR}]",
-        ":Hello ? world : nil" => "#{YELLOW}:#{CLEAR}#{YELLOW}Hello#{CLEAR} ? world : #{CYAN}#{BOLD}nil#{CLEAR}",
-        'raise "foo#{bar}baz"' => "raise #{RED}#{BOLD}\"#{CLEAR}#{RED}foo#{CLEAR}#{RED}\#{#{CLEAR}bar#{RED}}#{CLEAR}#{RED}baz#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}",
-        '["#{obj.inspect}"]' => "[#{RED}#{BOLD}\"#{CLEAR}#{RED}\#{#{CLEAR}obj.inspect#{RED}}#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}]",
-        'URI.parse "#{}"' => "#{BLUE}#{BOLD}#{UNDERLINE}URI#{CLEAR}.parse #{RED}#{BOLD}\"#{CLEAR}#{RED}\#{#{CLEAR}#{RED}}#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}",
+        ":Hello ? world : nil" => "#{YELLOW}:#{CLEAR}#{YELLOW}Hello#{CLEAR} ? #{CYAN}world#{CLEAR} : #{CYAN}#{BOLD}nil#{CLEAR}",
+        'raise "foo#{bar}baz"' => "#{CYAN}raise#{CLEAR} #{RED}#{BOLD}\"#{CLEAR}#{RED}foo#{CLEAR}#{RED}\#{#{CLEAR}#{CYAN}bar#{CLEAR}#{RED}}#{CLEAR}#{RED}baz#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}",
+        '["#{obj.inspect}"]' => "[#{RED}#{BOLD}\"#{CLEAR}#{RED}\#{#{CLEAR}#{CYAN}obj#{CLEAR}.#{CYAN}inspect#{CLEAR}#{RED}}#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}]",
+        'URI.parse "#{}"' => "#{BLUE}#{BOLD}#{UNDERLINE}URI#{CLEAR}.#{CYAN}parse#{CLEAR} #{RED}#{BOLD}\"#{CLEAR}#{RED}\#{#{CLEAR}#{RED}}#{CLEAR}#{RED}#{BOLD}\"#{CLEAR}",
         "begin\nrescue\nend" => "#{GREEN}begin#{CLEAR}\n#{GREEN}rescue#{CLEAR}\n#{GREEN}end#{CLEAR}",
-        "foo %w[bar]" => "foo #{RED}#{BOLD}%w[#{CLEAR}#{RED}bar#{CLEAR}#{RED}#{BOLD}]#{CLEAR}",
-        "foo %i[bar]" => "foo #{YELLOW}%i[#{CLEAR}#{YELLOW}bar#{CLEAR}#{YELLOW}]#{CLEAR}",
-        "foo :@bar, baz, :@@qux, :$quux" => "foo #{YELLOW}:#{CLEAR}#{YELLOW}@bar#{CLEAR}, baz, #{YELLOW}:#{CLEAR}#{YELLOW}@@qux#{CLEAR}, #{YELLOW}:#{CLEAR}#{YELLOW}$quux#{CLEAR}",
+        "foo %w[bar]" => "#{CYAN}foo#{CLEAR} #{RED}#{BOLD}%w[#{CLEAR}#{RED}bar#{CLEAR}#{RED}#{BOLD}]#{CLEAR}",
+        "foo %i[bar]" => "#{CYAN}foo#{CLEAR} #{YELLOW}%i[#{CLEAR}#{YELLOW}bar#{CLEAR}#{YELLOW}]#{CLEAR}",
+        "foo :@bar, baz, :@@qux, :$quux" => "#{CYAN}foo#{CLEAR} #{YELLOW}:#{CLEAR}#{YELLOW}@bar#{CLEAR}, #{CYAN}baz#{CLEAR}, #{YELLOW}:#{CLEAR}#{YELLOW}@@qux#{CLEAR}, #{YELLOW}:#{CLEAR}#{YELLOW}$quux#{CLEAR}",
         "`echo`" => "#{RED}#{BOLD}`#{CLEAR}#{RED}echo#{CLEAR}#{RED}#{BOLD}`#{CLEAR}",
         "\t" => Reline::Unicode.escape_for_print("\t") == '  ' ? '  ' : "\t", # not ^I
-        "foo(*%W(bar))" => "foo(*#{RED}#{BOLD}%W(#{CLEAR}#{RED}bar#{CLEAR}#{RED}#{BOLD})#{CLEAR})",
+        "foo(*%W(bar))" => "#{CYAN}foo#{CLEAR}(*#{RED}#{BOLD}%W(#{CLEAR}#{RED}bar#{CLEAR}#{RED}#{BOLD})#{CLEAR})",
         "$stdout" => "#{GREEN}#{BOLD}$stdout#{CLEAR}",
         "$&" => "#{GREEN}#{BOLD}$&#{CLEAR}",
         "$1" => "#{GREEN}#{BOLD}$1#{CLEAR}",
         "__END__" => "#{GREEN}__END__#{CLEAR}",
-        "foo\n__END__\nbar" => "foo\n#{GREEN}__END__#{CLEAR}\nbar",
-        "foo\n<<A\0\0bar\nA\nbaz" => "foo\n#{RED}<<A#{CLEAR}^@^@bar\n#{RED}A#{CLEAR}\nbaz",
+        "foo\n__END__\nbar" => "#{CYAN}foo#{CLEAR}\n#{GREEN}__END__#{CLEAR}\nbar",
+        "foo\n<<A\0\0bar\nA\nbaz" => "#{CYAN}foo#{CLEAR}\n#{RED}<<A#{CLEAR}^@^@bar\n#{RED}A#{CLEAR}\nbaz",
         "<<A+1\nA" => "#{RED}<<A#{CLEAR}+#{BLUE}#{BOLD}1#{CLEAR}\n#{RED}A#{CLEAR}",
         "4.5.6" => "#{MAGENTA}#{BOLD}4.5#{CLEAR}#{RED}#{REVERSE}.6#{CLEAR}",
-        "\e[0m\n" => "#{RED}#{REVERSE}^[#{CLEAR}[#{BLUE}#{BOLD}0#{CLEAR}m\n",
+        "\e[0m\n" => "#{RED}#{REVERSE}^[#{CLEAR}[#{BLUE}#{BOLD}0#{CLEAR}#{CYAN}m#{CLEAR}\n",
         "<<EOS\nhere\nEOS" => "#{RED}<<EOS#{CLEAR}\n#{RED}here#{CLEAR}\n#{RED}EOS#{CLEAR}",
         "[1]]]\u0013" => "[#{BLUE}#{BOLD}1#{CLEAR}]#{RED}#{REVERSE}]#{CLEAR}#{RED}#{REVERSE}]#{CLEAR}#{RED}#{REVERSE}^S#{CLEAR}",
-        "def req(true) end" => "#{GREEN}def#{CLEAR} #{BLUE}#{BOLD}req#{CLEAR}(#{RED}#{REVERSE}true#{CLEAR}#{RED}#{REVERSE})#{CLEAR} #{GREEN}end#{CLEAR}",
+        "def req(true) end" => "#{GREEN}def#{CLEAR} #{CYAN}#{BOLD}req#{CLEAR}(#{RED}#{REVERSE}true#{CLEAR}#{RED}#{REVERSE})#{CLEAR} #{GREEN}end#{CLEAR}",
         "nil = 1" => "#{CYAN}#{BOLD}nil#{CLEAR} #{RED}#{REVERSE}=#{CLEAR} #{BLUE}#{BOLD}1#{CLEAR}",
         "alias $x $1" => "#{GREEN}alias#{CLEAR} #{GREEN}#{BOLD}$x#{CLEAR} #{RED}#{REVERSE}$1#{CLEAR}",
         "class bad; end" => "#{GREEN}class#{CLEAR} #{RED}#{REVERSE}bad#{CLEAR}; #{GREEN}end#{CLEAR}",
-        "def req(@a) end" => "#{GREEN}def#{CLEAR} #{BLUE}#{BOLD}req#{CLEAR}(#{RED}#{REVERSE}@a#{CLEAR}) #{GREEN}end#{CLEAR}",
+        "def req(@a) end" => "#{GREEN}def#{CLEAR} #{CYAN}#{BOLD}req#{CLEAR}(#{RED}#{REVERSE}@a#{CLEAR}) #{GREEN}end#{CLEAR}",
+        "a.foo" => "#{CYAN}a#{CLEAR}.#{CYAN}foo#{CLEAR}",
+        "a.foo = 1" => "#{CYAN}a#{CLEAR}.#{CYAN}foo#{CLEAR} = #{BLUE}#{BOLD}1#{CLEAR}",
+        "a.foo += 1" => "#{CYAN}a#{CLEAR}.#{CYAN}foo#{CLEAR} += #{BLUE}#{BOLD}1#{CLEAR}",
+        "a.foo &&= 1" => "#{CYAN}a#{CLEAR}.#{CYAN}foo#{CLEAR} &&= #{BLUE}#{BOLD}1#{CLEAR}",
+        "a.foo ||= 1" => "#{CYAN}a#{CLEAR}.#{CYAN}foo#{CLEAR} ||= #{BLUE}#{BOLD}1#{CLEAR}",
+        "a[:foo]" => "#{CYAN}a#{CLEAR}[#{YELLOW}:#{CLEAR}#{YELLOW}foo#{CLEAR}]",
+        "a[:foo] = 1" => "#{CYAN}a#{CLEAR}[#{YELLOW}:#{CLEAR}#{YELLOW}foo#{CLEAR}] = #{BLUE}#{BOLD}1#{CLEAR}",
+        "a+1" => "#{CYAN}a#{CLEAR}+#{BLUE}#{BOLD}1#{CLEAR}",
+        "a.+(1)" => "#{CYAN}a#{CLEAR}.#{CYAN}+#{CLEAR}(#{BLUE}#{BOLD}1#{CLEAR})",
+        "-a" => "-#{CYAN}a#{CLEAR}",
+        "a.-@" => "#{CYAN}a#{CLEAR}.#{CYAN}-@#{CLEAR}",
+        'Foo(1)' => "#{BLUE}#{BOLD}#{UNDERLINE}Foo#{CLEAR}(#{BLUE}#{BOLD}1#{CLEAR})",
+        'Foo::Bar(1)' => "#{BLUE}#{BOLD}#{UNDERLINE}Foo#{CLEAR}::#{BLUE}#{BOLD}#{UNDERLINE}Bar#{CLEAR}(#{BLUE}#{BOLD}1#{CLEAR})",
+        'Foo.Bar(1)' => "#{BLUE}#{BOLD}#{UNDERLINE}Foo#{CLEAR}.#{CYAN}Bar#{CLEAR}(#{BLUE}#{BOLD}1#{CLEAR})",
+        'Foo&.Bar(1)' => "#{BLUE}#{BOLD}#{UNDERLINE}Foo#{CLEAR}&.#{CYAN}Bar#{CLEAR}(#{BLUE}#{BOLD}1#{CLEAR})",
+        'ENV' => "#{BLUE}#{BOLD}#{UNDERLINE}ENV#{CLEAR}",
       }
 
       tests.each do |code, result|
@@ -138,9 +156,9 @@ def test_colorize_code
 
     def test_colorize_code_with_local_variables
       code = "a /(b +1)/i"
-      result_without_lvars = "a #{RED}#{BOLD}/#{CLEAR}#{RED}(b +1)#{CLEAR}#{RED}#{BOLD}/i#{CLEAR}"
-      result_with_lvar = "a /(b #{BLUE}#{BOLD}+1#{CLEAR})/i"
-      result_with_lvars = "a /(b +#{BLUE}#{BOLD}1#{CLEAR})/i"
+      result_without_lvars = "#{CYAN}a#{CLEAR} #{RED}#{BOLD}/#{CLEAR}#{RED}(b +1)#{CLEAR}#{RED}#{BOLD}/i#{CLEAR}"
+      result_with_lvar = "a /(#{CYAN}b#{CLEAR} #{BLUE}#{BOLD}+1#{CLEAR})/#{CYAN}i#{CLEAR}"
+      result_with_lvars = "a /(b +#{BLUE}#{BOLD}1#{CLEAR})/#{CYAN}i#{CLEAR}"
 
       assert_equal_with_term(result_without_lvars, code)
       assert_equal_with_term(result_with_lvar, code, local_variables: [:a])
@@ -169,7 +187,7 @@ def test_colorize_code_complete_false
         "'foo' + 'bar" => "#{RED}#{BOLD}'#{CLEAR}#{RED}foo#{CLEAR}#{RED}#{BOLD}'#{CLEAR} + #{RED}#{BOLD}'#{CLEAR}#{RED}bar#{CLEAR}",
         "('foo" => "(#{RED}#{BOLD}'#{CLEAR}#{RED}foo#{CLEAR}",
         "if true" => "#{GREEN}if#{CLEAR} #{CYAN}#{BOLD}true#{CLEAR}",
-        "tap do end end tap do" => "tap #{GREEN}do#{CLEAR} #{GREEN}end#{CLEAR} #{RED}#{REVERSE}end#{CLEAR} tap #{GREEN}do#{CLEAR}",
+        "tap do end end tap do" => "#{CYAN}tap#{CLEAR} #{GREEN}do#{CLEAR} #{GREEN}end#{CLEAR} #{RED}#{REVERSE}end#{CLEAR} #{CYAN}tap#{CLEAR} #{GREEN}do#{CLEAR}",
 
         # Specially handled cases
         "class" => "#{GREEN}class#{CLEAR}",
@@ -178,17 +196,17 @@ def test_colorize_code_complete_false
         "module" => "#{GREEN}module#{CLEAR}",
         "module#" => "#{GREEN}module#{CLEAR}#{BLUE}#{BOLD}\##{CLEAR}",
         "module;" => "#{RED}#{REVERSE}module#{CLEAR};",
-        "class owner_module" => "#{GREEN}class#{CLEAR} owner_module",
-        "module owner_module" => "#{GREEN}module#{CLEAR} owner_module",
+        "class owner_module" => "#{GREEN}class#{CLEAR} #{CYAN}owner_module#{CLEAR}",
+        "module owner_module" => "#{GREEN}module#{CLEAR} #{CYAN}owner_module#{CLEAR}",
         "a, b" => "a, b",
         "a, b#" => "a, b#{BLUE}#{BOLD}\##{CLEAR}",
         "a, b;" => "#{RED}#{REVERSE}a, b#{CLEAR};",
-        "def f(a,#" => "#{GREEN}def#{CLEAR} #{BLUE}#{BOLD}f#{CLEAR}(a,#{BLUE}#{BOLD}\##{CLEAR}",
+        "def f(a,#" => "#{GREEN}def#{CLEAR} #{CYAN}#{BOLD}f#{CLEAR}(a,#{BLUE}#{BOLD}\##{CLEAR}",
         "[*" => "[*",
-        "f(*" => "f(*",
-        "f(**" => "f(**",
-        "f(&" => "f(&",
-        "def f =" => "#{GREEN}def#{CLEAR} #{BLUE}#{BOLD}f#{CLEAR} =",
+        "f(*" => "#{CYAN}f#{CLEAR}(*",
+        "f(**" => "#{CYAN}f#{CLEAR}(**",
+        "f(&" => "#{CYAN}f#{CLEAR}(&",
+        "def f =" => "#{GREEN}def#{CLEAR} #{CYAN}#{BOLD}f#{CLEAR} =",
         "=begin" => "#{BLUE}#{BOLD}=begin#{CLEAR}",
       }.each do |code, result|
         assert_equal_with_term(result, code, complete: false)