Colorize IRB's code_around_binding

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

1 files changed, 83 insertions, 0 deletions

diff --git a/lib/irb/color.rb b/lib/irb/color.rb
new file mode 100644
index 0000000000..150e9e5666
--- /dev/null
+++ b/lib/irb/color.rb
@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+require 'ripper'
+
+module IRB # :nodoc:
+  module Color
+    CLEAR     = 0
+    BOLD      = 1
+    UNDERLINE = 4
+    RED       = 31
+    GREEN     = 32
+    BLUE      = 34
+    MAGENTA   = 35
+    CYAN      = 36
+
+    TOKEN_KEYWORDS = {
+      on_kw: ['nil', 'self', 'true', 'false'],
+      on_const: ['ENV'],
+    }
+
+    TOKEN_SEQ_EXPRS = {
+      on_CHAR:            [[BLUE, BOLD],            [Ripper::EXPR_END]],
+      on_const:           [[BLUE, BOLD, UNDERLINE], [Ripper::EXPR_ARG, Ripper::EXPR_CMDARG]],
+      on_embexpr_beg:     [[RED],                   [Ripper::EXPR_BEG, Ripper::EXPR_END]],
+      on_embexpr_end:     [[RED],                   [Ripper::EXPR_END, Ripper::EXPR_ENDFN]],
+      on_ident:           [[BLUE, BOLD],            [Ripper::EXPR_ENDFN]],
+      on_int:             [[BLUE, BOLD],            [Ripper::EXPR_END]],
+      on_float:           [[MAGENTA, BOLD],         [Ripper::EXPR_END]],
+      on_kw:              [[GREEN],                 [Ripper::EXPR_CLASS, Ripper::EXPR_BEG, Ripper::EXPR_END, Ripper::EXPR_FNAME]],
+      on_label:           [[MAGENTA],               [Ripper::EXPR_LABELED]],
+      on_qwords_beg:      [[RED],                   [Ripper::EXPR_BEG]],
+      on_regexp_beg:      [[RED],                   [Ripper::EXPR_BEG]],
+      on_regexp_end:      [[RED],                   [Ripper::EXPR_BEG]],
+      on_symbeg:          [[BLUE, BOLD],            [Ripper::EXPR_FNAME]],
+      on_tstring_beg:     [[RED],                   [Ripper::EXPR_BEG, Ripper::EXPR_END, Ripper::EXPR_ARG, Ripper::EXPR_CMDARG]],
+      on_tstring_content: [[RED],                   [Ripper::EXPR_BEG, Ripper::EXPR_ARG, Ripper::EXPR_CMDARG]],
+      on_tstring_end:     [[RED],                   [Ripper::EXPR_END]],
+    }
+
+    class << self
+      def colorable?
+        $stdout.tty? && ENV.key?('TERM')
+      end
+
+      def clear
+        return '' unless colorable?
+        "\e[#{CLEAR}m"
+      end
+
+      def colorize(text, seq)
+        return text unless colorable?
+        "#{seq.map { |s| "\e[#{const_get(s)}m" }.join('')}#{text}#{clear}"
+      end
+
+      def colorize_code(code)
+        return code unless colorable?
+
+        colored = +''
+        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}"
+          else
+            colored << str
+          end
+        end
+        colored
+      end
+
+      private
+
+      def dispatch_seq(token, expr, str)
+        if token == :on_comment
+          [BLUE, BOLD]
+        elsif TOKEN_KEYWORDS.fetch(token, []).include?(str)
+          [CYAN, BOLD]
+        elsif (seq, exprs = TOKEN_SEQ_EXPRS[token]; exprs&.any? { |e| (expr & e) != Ripper::EXPR_NONE })
+          seq
+        else
+          nil
+        end
+      end
+    end
+  end
+end