summaryrefslogtreecommitdiff
path: root/lib/irb/color.rb
blob: 71557ffbb92849cf60ed2afb40454e9bbd6b6efc (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
# frozen_string_literal: true
require 'reline'
require 'ripper'

module IRB # :nodoc:
  module Color
    CLEAR     = 0
    BOLD      = 1
    UNDERLINE = 4
    REVERSE   = 7
    RED       = 31
    GREEN     = 32
    YELLOW    = 33
    BLUE      = 34
    MAGENTA   = 35
    CYAN      = 36

    TOKEN_KEYWORDS = {
      on_kw: ['nil', 'self', 'true', 'false', '__FILE__', '__LINE__'],
      on_const: ['ENV'],
    }
    private_constant :TOKEN_KEYWORDS

    # A constant of all-bit 1 to match any Ripper's state in #dispatch_seq
    ALL = -1
    private_constant :ALL

    begin
      # Following pry's colors where possible, but sometimes having a compromise like making
      # backtick and regexp as red (string's color, because they're sharing tokens).
      TOKEN_SEQ_EXPRS = {
        on_CHAR:            [[BLUE, BOLD],            ALL],
        on_backtick:        [[RED],                   ALL],
        on_comment:         [[BLUE, BOLD],            ALL],
        on_const:           [[BLUE, BOLD, UNDERLINE], ALL],
        on_embexpr_beg:     [[RED],                   ALL],
        on_embexpr_end:     [[RED],                   ALL],
        on_embvar:          [[RED],                   ALL],
        on_float:           [[MAGENTA, BOLD],         ALL],
        on_gvar:            [[GREEN, BOLD],           ALL],
        on_heredoc_beg:     [[RED],                   ALL],
        on_heredoc_end:     [[RED],                   ALL],
        on_ident:           [[BLUE, BOLD],            Ripper::EXPR_ENDFN],
        on_imaginary:       [[BLUE, BOLD],            ALL],
        on_int:             [[BLUE, BOLD],            ALL],
        on_kw:              [[GREEN],                 ALL],
        on_label:           [[MAGENTA],               ALL],
        on_label_end:       [[RED],                   ALL],
        on_qsymbols_beg:    [[RED],                   ALL],
        on_qwords_beg:      [[RED],                   ALL],
        on_rational:        [[BLUE, BOLD],            ALL],
        on_regexp_beg:      [[RED, BOLD],             ALL],
        on_regexp_end:      [[RED, BOLD],             ALL],
        on_symbeg:          [[YELLOW],                ALL],
        on_tstring_beg:     [[RED],                   ALL],
        on_tstring_content: [[RED],                   ALL],
        on_tstring_end:     [[RED],                   ALL],
        on_words_beg:       [[RED],                   ALL],
        on_parse_error:     [[RED, REVERSE],          ALL],
        compile_error:      [[RED, REVERSE],          ALL],
      }
    rescue NameError
      # Give up highlighting Ripper-incompatible older Ruby
      TOKEN_SEQ_EXPRS = {}
    end
    private_constant :TOKEN_SEQ_EXPRS

    class << self
      def colorable?
        $stdout.tty? && (/mswin|mingw/ =~ RUBY_PLATFORM || (ENV.key?('TERM') && ENV['TERM'] != 'dumb'))
      end

      def inspect_colorable?(obj)
        case obj
        when String, Symbol, Regexp, Integer, Float, FalseClass, TrueClass, NilClass
          true
        when Hash
          obj.all? { |k, v| inspect_colorable?(k) && inspect_colorable?(v) }
        when Array
          obj.all? { |o| inspect_colorable?(o) }
        when Range
          inspect_colorable?(obj.begin) && inspect_colorable?(obj.end)
        when Module
          !obj.name.nil?
        else
          false
        end
      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 scan(code)
        Ripper::Lexer.new(code).scan
      end

      def colorize_code(code)
        return code unless colorable?

        symbol_state = SymbolState.new
        colored = +''
        length = 0

        scan(code).each do |elem|
          token = elem.event
          str = elem.tok
          expr = elem.state
          in_symbol = symbol_state.scan_token(token)
          if seq = dispatch_seq(token, expr, str, in_symbol: in_symbol)
            Reline::Unicode.escape_for_print(str).each_line do |line|
              colored << seq.map { |s| "\e[#{s}m" }.join('')
              colored << line.sub(/\Z/, clear)
            end
          else
            colored << Reline::Unicode.escape_for_print(str)
          end
          length += str.length
        end

        # give up colorizing incomplete Ripper tokens
        return code if length != code.length

        colored
      end

      private

      def dispatch_seq(token, expr, str, in_symbol:)
        if token == :on_parse_error or token == :compile_error
          TOKEN_SEQ_EXPRS[token][0]
        elsif in_symbol
          [YELLOW]
        elsif TOKEN_KEYWORDS.fetch(token, []).include?(str)
          [CYAN, BOLD]
        elsif (seq, exprs = TOKEN_SEQ_EXPRS[token]; (expr & (exprs || 0)) != 0)
          seq
        else
          nil
        end
      end
    end

    # A class to manage a state to know whether the current token is for Symbol or not.
    class SymbolState
      def initialize
        # Push `true` to detect Symbol. `false` to increase the nest level for non-Symbol.
        @stack = []
      end

      # Return true if the token is a part of Symbol.
      def scan_token(token)
        prev_state = @stack.last
        case token
        when :on_symbeg
          @stack << true
        when :on_ident, :on_op, :on_const, :on_ivar, :on_cvar, :on_gvar, :on_kw
          if @stack.last # Pop only when it's Symbol
            @stack.pop
            return prev_state
          end
        when :on_tstring_beg
          @stack << false
        when :on_embexpr_beg
          @stack << false
          return prev_state
        when :on_tstring_end # :on_tstring_end may close Symbol
          @stack.pop
          return prev_state
        when :on_embexpr_end
          @stack.pop
        end
        @stack.last
      end
    end
    private_constant :SymbolState
  end
end