/extensions/wirble/lib/wirble.rb

  1require 'ostruct'
  2
  3#
  4# Wirble: A collection of useful Irb features.
  5#
  6# To use, add the following to your ~/.irbrc:
  7#
  8#   require 'rubygems'
  9#   require 'wirble'
 10#   Wirble.init
 11#
 12# If you want color in Irb, add this to your ~/.irbrc as well:
 13#
 14#   Wirble.colorize
 15#
 16# Note:  I spent a fair amount of time documenting this code in the
 17# README.  If you've installed via RubyGems, root around your cache a
 18# little bit (or fire up gem_server) and read it before you tear your
 19# hair out sifting through the code below.
 20# 
 21module Wirble
 22  VERSION = '0.1.3.2'
 23
 24  #
 25  # Load internal Ruby features, including pp, tab-completion, 
 26  # and a simple prompt.
 27  #
 28  module Internals
 29    # list of internal libraries to automatically load
 30    LIBRARIES = %w{pp irb/completion}
 31
 32    #
 33    # load libraries
 34    #
 35    def self.init_libraries
 36      LIBRARIES.each do |lib| 
 37        begin
 38          require lib 
 39        rescue LoadError
 40          nil
 41        end
 42      end
 43    end
 44
 45    #
 46    # Set a simple prompt, unless a custom one has been specified.
 47    #
 48    def self.init_prompt
 49      # set the prompt
 50      if IRB.conf[:PROMPT_MODE] == :DEFAULT
 51        IRB.conf[:PROMPT_MODE] = :SIMPLE
 52      end
 53    end
 54
 55    #
 56    # Load all Ruby internal features.
 57    #
 58    def self.init(opt = nil)
 59      init_libraries unless opt && opt[:skip_libraries]
 60      init_prompt unless opt && opt[:skip_prompt]
 61    end
 62  end
 63
 64  #
 65  # Basic IRB history support.  This is based on the tips from 
 66  # http://wiki.rubygarden.org/Ruby/page/show/Irb/TipsAndTricks
 67  #
 68  class History
 69    DEFAULTS = {
 70      :history_path   => ENV['IRB_HISTORY_FILE'] || "~/.irb_history",
 71      :history_size   => (ENV['IRB_HISTORY_SIZE'] || 1000).to_i,
 72      :history_perms  => File::WRONLY | File::CREAT | File::TRUNC,
 73      :history_uniq   => true,
 74    }
 75 
 76    private
 77
 78    def say(*args)
 79      puts(*args) if @verbose
 80    end
 81
 82    def cfg(key)
 83      @opt["history_#{key}".intern]
 84    end
 85
 86    def save_history
 87      path, max_size, perms, uniq = %w{path size perms uniq}.map { |v| cfg(v) }
 88
 89      # read lines from history, and truncate the list (if necessary)
 90      lines = Readline::HISTORY.to_a
 91
 92      lines.reverse! if reverse = uniq.to_s == 'reverse'
 93      lines.uniq!    if uniq
 94      lines.reverse! if reverse
 95
 96      lines.slice!(0, lines.size - max_size) if lines.size > max_size
 97
 98      # write the history file
 99      real_path = File.expand_path(path)
100      File.open(real_path, perms) { |fh| fh.puts lines }
101      say 'Saved %d lines to history file %s.' % [lines.size, path]
102    end
103
104    def load_history
105      # expand history file and make sure it exists
106      real_path = File.expand_path(cfg('path'))
107      unless File.exist?(real_path)
108        say "History file #{real_path} doesn't exist."
109        return
110      end
111
112      # read lines from file and add them to history
113      lines = File.readlines(real_path).map { |line| line.chomp }
114      Readline::HISTORY.push(*lines)
115
116      say 'Read %d lines from history file %s' % [lines.size, cfg('path')]
117    end
118
119    public
120
121    def initialize(opt = nil)
122      @opt = DEFAULTS.merge(opt || {})
123      return unless defined? Readline::HISTORY
124      load_history
125      Kernel.at_exit { save_history }
126    end
127  end
128
129  #
130  # Add color support to IRB.
131  #
132  module Colorize
133    #
134    # Tokenize an inspection string.
135    #
136    module Tokenizer
137      def self.tokenize(str)
138        raise 'missing block' unless block_given?
139        chars = str.split(//)
140
141        # $stderr.puts "DEBUG: chars = #{chars.join(',')}"
142
143        state, val, i, lc = [], '', 0, nil
144        while i <= chars.size
145          repeat = false
146          c = chars[i]
147
148          # $stderr.puts "DEBUG: state = #{state}"
149
150          case state[-1]
151          when nil
152            case c
153            when ':'
154              state << :symbol
155            when '"'
156              state << :string
157            when '#'
158              state << :object
159            when /[a-z]/i
160              state << :keyword
161              repeat = true
162            when /[0-9-]/
163              state << :number
164              repeat = true
165            when '{'
166              yield :open_hash, '{'
167            when '['
168              yield :open_array, '['
169            when ']'
170              yield :close_array, ']'
171            when '}'
172              yield :close_hash, '}'
173            when /\s/
174              yield :whitespace, c
175            when ','
176              yield :comma, ','
177            when '>'
178              yield :refers, '=>' if lc == '='
179            when '.'
180              yield :range, '..' if lc == '.'
181            when '='
182              # ignore these, they're used elsewhere
183              nil
184            else 
185              # $stderr.puts "DEBUG: ignoring char #{c}"
186            end
187          when :symbol
188            case c
189            # XXX: should have =, but that messes up foo=>bar
190            when /[a-z0-9_!?]/
191              val << c
192            else
193              yield :symbol_prefix, ':'
194              yield state[-1], val
195              state.pop; val = ''
196              repeat = true
197            end
198          when :string
199            case c
200            when '"'
201              if lc == "\\"
202                val[-1] = ?"
203              else
204                yield :open_string, '"'
205                yield state[-1], val
206                state.pop; val = ''
207                yield :close_string, '"'
208              end
209            else
210              val << c
211            end
212          when :keyword
213            case c
214            when /[a-z0-9_]/i
215              val << c
216            else
217              # is this a class?
218              st = val =~ /^[A-Z]/ ? :class : state[-1]
219
220              yield st, val
221              state.pop; val = ''
222              repeat = true
223            end
224          when :number
225            case c
226            when /[0-9e-]/
227              val << c
228            when '.'
229              if lc == '.'
230                val[/\.$/] = ''
231                yield state[-1], val
232                state.pop; val = ''
233                yield :range, '..'
234              else
235                val << c
236              end
237            else
238              yield state[-1], val
239              state.pop; val = ''
240              repeat = true
241            end
242          when :object
243            case c
244            when '<' 
245              yield :open_object, '#<'
246              state << :object_class
247            when ':' 
248              state << :object_addr
249            when '@' 
250              state << :object_line
251            when '>'
252              yield :close_object, '>'
253              state.pop; val = ''
254            end
255          when :object_class
256            case c
257            when ':'
258              yield state[-1], val
259              state.pop; val = ''
260              repeat = true
261            else
262              val << c
263            end
264          when :object_addr
265            case c
266            when '>'
267            when '@'
268              yield :object_addr_prefix, ':'
269              yield state[-1], val
270              state.pop; val = ''
271              repeat = true
272            else
273              val << c
274            end
275          when :object_line
276            case c
277            when '>'
278              yield :object_line_prefix, '@'
279              yield state[-1], val
280              state.pop; val = ''
281              repeat = true
282            else
283              val << c
284            end
285          else
286            raise "unknown state #{state}"
287          end
288
289          unless repeat
290            i += 1
291            lc = c
292          end
293        end
294      end
295    end
296
297    #
298    # Terminal escape codes for colors.
299    #
300    module Color
301      COLORS = {
302        :nothing      => '0;0',
303        :black        => '0;30',
304        :red          => '0;31',
305        :green        => '0;32',
306        :brown        => '0;33',
307        :blue         => '0;34',
308        :cyan         => '0;36',
309        :purple       => '0;35',
310        :light_gray   => '0;37',
311        :dark_gray    => '1;30',
312        :light_red    => '1;31',
313        :light_green  => '1;32',
314        :yellow       => '1;33',
315        :light_blue   => '1;34',
316        :light_cyan   => '1;36',
317        :light_purple => '1;35',
318        :white        => '1;37',
319      }
320      
321      #
322      # Return the escape code for a given color.
323      #
324      def self.escape(key)
325        COLORS.key?(key) && "\033[#{COLORS[key]}m"
326      end
327    end
328
329    #
330    # Default Wirble color scheme.
331    # 
332    DEFAULT_COLORS = {
333      # delimiter colors
334      :comma              => :blue,
335      :refers             => :blue,
336
337      # container colors (hash and array)
338      :open_hash          => :green,
339      :close_hash         => :green,
340      :open_array         => :green,
341      :close_array        => :green,
342
343      # object colors
344      :open_object        => :light_red,
345      :object_class       => :white,
346      :object_addr_prefix => :blue,
347      :object_line_prefix => :blue,
348      :close_object       => :light_red,
349
350      # symbol colors
351      :symbol             => :yellow,
352      :symbol_prefix      => :yellow,
353
354      # string colors
355      :open_string        => :red,
356      :string             => :cyan,
357      :close_string       => :red,
358
359      # misc colors
360      :number             => :cyan,
361      :keyword            => :green,
362      :class              => :light_green,
363      :range              => :red,
364    }
365
366    #
367    # Fruity testing colors.
368    # 
369    TESTING_COLORS = {
370      :comma            => :red,
371      :refers           => :red,
372      :open_hash        => :blue,
373      :close_hash       => :blue,
374      :open_array       => :green,
375      :close_array      => :green,
376      :open_object      => :light_red,
377      :object_class     => :light_green,
378      :object_addr      => :purple,
379      :object_line      => :light_purple,
380      :close_object     => :light_red,
381      :symbol           => :yellow,
382      :symbol_prefix    => :yellow,
383      :number           => :cyan,
384      :string           => :cyan,
385      :keyword          => :white,
386      :range            => :light_blue,
387    }
388
389    #
390    # Set color map to hash
391    # 
392    def self.colors=(hash)
393      @colors = hash
394    end
395
396    #
397    # Get current color map
398    # 
399    def self.colors
400      @colors ||= {}.update(DEFAULT_COLORS)
401    end
402
403    #
404    # Return a string with the given color.
405    #
406    def self.colorize_string(str, color)
407      col, nocol = [color, :nothing].map { |key| Color.escape(key) }
408      col ? "#{col}#{str}#{nocol}" : str
409    end
410
411    #
412    # Colorize the results of inspect
413    # 
414    def self.colorize(str)
415      begin
416        ret, nocol = '', Color.escape(:nothing)
417        Tokenizer.tokenize(str) do |tok, val|
418          # c = Color.escape(colors[tok])
419          ret << colorize_string(val, colors[tok])
420        end
421        ret
422      rescue
423        # catch any errors from the tokenizer (just in case)
424        str
425      end
426    end
427
428    #
429    # Enable colorized IRB results.
430    # 
431    def self.enable(custom_colors = nil)
432      # if there's a better way to do this, I'm all ears.
433      ::IRB::Irb.class_eval do
434        alias :non_color_output_value  :output_value
435
436        def output_value
437          if @context.inspect?
438            val = Colorize.colorize(@context.last_value.inspect)
439            p val
440            printf @context.return_format, val
441          else
442            printf @context.return_format, @context.last_value
443          end
444        end
445      end
446
447      self.colors = custom_colors if custom_colors
448    end
449
450    #
451    # Disable colorized IRB results.
452    # 
453    def self.disable
454      ::IRB::Irb.class_eval do
455        alias :output_value  :non_color_output_value
456      end
457    end
458  end
459
460  #
461  # Convenient shortcut methods.
462  #
463  module Shortcuts
464    #
465    # Print object methods, sorted by name. (excluding methods that
466    # exist in the class Object) .
467    #
468    def po(o)
469      o.methods.sort - Object.methods
470    end
471
472    #
473    # Print object constants, sorted by name.
474    #
475    def poc(o)
476      o.constants.sort
477    end
478  end
479
480  #
481  # Convenient shortcut for ri
482  #
483  module RiShortcut
484    def self.init
485      Kernel.class_eval {
486        def ri(arg)
487           puts `ri '#{arg}'`
488        end
489      }
490
491      Module.instance_eval {
492         def ri(meth=nil)
493           if meth
494             if instance_methods(false).include? meth.to_s
495               puts `ri #{self}##{meth}`
496             else
497               super
498             end
499           else
500             puts `ri #{self}`
501           end
502         end
503      }
504    end
505  end
506
507
508
509  #
510  # Enable color results.
511  #
512  def self.colorize(custom_colors = nil)
513    Colorize.enable(custom_colors)
514  end
515
516  #
517  # Load everything except color.
518  #
519  def self.init(opt = nil)
520    # make sure opt isn't nil
521    opt ||= {}
522
523    # load internal irb/ruby features
524    Internals.init(opt) unless opt && opt[:skip_internals]
525
526    # load the history
527    History.new(opt) unless opt && opt[:skip_history]
528
529    # load shortcuts
530    unless opt && opt[:skip_shortcuts]
531      # load ri shortcuts
532      RiShortcut.init
533
534      # include common shortcuts
535      Object.class_eval { include Shortcuts }
536    end
537
538    colorize(opt[:colors]) if opt && opt[:init_colors]
539  end
540end
541