Merged in capistrano_colors gem, and renamed to 'log_formatters', since it does much more than just colors

lib/capistrano/configuration.rb

@@ -5,6 +5,7 @@ require 'capistrano/configuration/callbacks'
 require 'capistrano/configuration/connections'
 require 'capistrano/configuration/execution'
 require 'capistrano/configuration/loading'
+require 'capistrano/configuration/log_formatters'
 require 'capistrano/configuration/namespaces'
 require 'capistrano/configuration/roles'
 require 'capistrano/configuration/servers'
@@ -34,7 +35,7 @@ module Capistrano
 
     # The includes must come at the bottom, since they may redefine methods
     # defined in the base class.
-    include AliasTask, Connections, Execution, Loading, Namespaces, Roles, Servers, Variables
+    include AliasTask, Connections, Execution, Loading, LogFormatters, Namespaces, Roles, Servers, Variables
 
     # Mix in the actions
     include Actions::FileTransfer, Actions::Inspect, Actions::Invocation

lib/capistrano/configuration/log_formatters.rb

@@ -0,0 +1,71 @@
+# Add custom log formatters
+#
+# Passing a hash or a array of hashes with custom log formatters.
+#
+# Add the following to your deploy.rb or in your ~/.caprc
+#
+# == Example:
+#
+#   capistrano_log_formatters = [
+#     { :match => /command finished/,       :color => :hide,      :priority => 10, :prepend => "$$$" },
+#     { :match => /executing command/,      :color => :blue,      :priority => 10, :style => :underscore, :timestamp => true },
+#     { :match => /^transaction: commit$/,  :color => :magenta,   :priority => 10, :style => :blink },
+#     { :match => /git/,                    :color => :white,     :priority => 20, :style => :reverse }
+#   ]
+#
+#   format_logs capistrano_log_formatters
+#
+# You can call format_logs multiple times, with either a hash or an array of hashes.
+#
+# == Colors:
+#
+# :color can have the following values:
+#
+# * :hide  (hides the row completely)
+# * :none
+# * :black
+# * :red
+# * :green
+# * :yellow
+# * :blue
+# * :magenta
+# * :cyan
+# * :white
+#
+# == Styles:
+#
+# :style can have the following values:
+#
+# * :bright
+# * :dim
+# * :underscore
+# * :blink
+# * :reverse
+# * :hidden
+#
+#
+#  == Text alterations
+#
+# :prepend gives static text to be prepended to the output
+# :replace replaces the matched text in the output
+# :timestamp adds the current time before the output
+
+module Capistrano
+  class Configuration
+    module LogFormatters
+      def log_formatter(options)
+        if options.class == Array
+          options.each do |option|
+            Capistrano::Logger.add_formatter(option)
+          end
+        else
+          Capistrano::Logger.add_formatter(options)
+        end
+      end
+
+      def disable_log_formatters
+        @logger.disable_formatters = true
+      end
+    end
+  end
+end

lib/capistrano/logger.rb

@@ -1,7 +1,6 @@
 module Capistrano
   class Logger #:nodoc:
-    attr_accessor :level
-    attr_reader   :device
+    attr_accessor :level, :device, :disable_formatters
 
     IMPORTANT = 0
     INFO      = 1
@@ -10,6 +9,59 @@ module Capistrano
 
     MAX_LEVEL = 3
 
+    COLORS = {
+      :none     => "0",
+      :black    => "30",
+      :red      => "31",
+      :green    => "32",
+      :yellow   => "33",
+      :blue     => "34",
+      :magenta  => "35",
+      :cyan     => "36",
+      :white    => "37"
+    }
+
+    STYLES = {
+      :bright     => 1,
+      :dim        => 2,
+      :underscore => 4,
+      :blink      => 5,
+      :reverse    => 7,
+      :hidden     => 8
+    }
+
+    # Set up default formatters
+    @formatters = [
+      # TRACE
+      { :match => /command finished/,          :color => :white,   :style => :dim, :level => 3, :priority => -10 },
+      { :match => /executing locally/,         :color => :yellow,  :level => 3, :priority => -20 },
+
+      # DEBUG
+      { :match => /executing `.*/,             :color => :green,   :level => 2, :priority => -10, :timestamp => true },
+      { :match => /.*/,                        :color => :yellow,  :level => 2, :priority => -30 },
+
+      # INFO
+      { :match => /.*out\] (fatal:|ERROR:).*/, :color => :red,     :level => 1, :priority => -10 },
+      { :match => /Permission denied/,         :color => :red,     :level => 1, :priority => -20 },
+      { :match => /sh: .+: command not found/, :color => :magenta, :level => 1, :priority => -30 },
+
+      # IMPORTANT
+      { :match => /^err ::/,                   :color => :red,     :level => 0, :priority => -10 },
+      { :match => /.*/,                        :color => :blue,    :level => 0, :priority => -20 }
+    ]
+
+    class << self
+      def add_formatter(options) #:nodoc:
+        @formatters.push(options)
+        @sorted_formatters = nil
+      end
+
+      def sorted_formatters
+        # Sort matchers in reverse order so we can break if we found a match.
+        @sorted_formatters ||= @formatters.sort_by { |i| -(i[:priority] || i[:prio] || 0) }
+      end
+    end
+
     def initialize(options={})
       output = options[:output] || $stderr
       if output.respond_to?(:puts)
@@ -21,6 +73,7 @@ module Capistrano
 
       @options = options
       @level = options[:level] || 0
+      @disable_formatters = options[:disable_formatters]
     end
 
     def close
@@ -29,6 +82,42 @@ module Capistrano
 
     def log(level, message, line_prefix=nil)
       if level <= self.level
+        # Only format output if device is a TTY or formatters are not disabled
+        if device.tty? && !@disable_formatters
+          color = :none
+          style = nil
+
+          Logger.sorted_formatters.each do |formatter|
+            if (formatter[:level] == level || formatter[:level].nil?)
+              if message =~ formatter[:match] || line_prefix =~ formatter[:match]
+                color = formatter[:color] if formatter[:color]
+                style = formatter[:style] || formatter[:attribute] # (support original cap colors)
+                message.gsub!(formatter[:match], formatter[:replace]) if formatter[:replace]
+                message = formatter[:prepend] + message unless formatter[:prepend].nil?
+                message = message + formatter[:append] unless formatter[:append].nil?
+                message = Time.now.strftime('%Y-%m-%d %T') + ' ' + message if formatter[:timestamp]
+                break unless formatter[:replace]
+              end
+            end
+          end
+
+          if color == :hide
+            # Don't do anything if color is set to :hide
+            return false
+          end
+
+          term_color = COLORS[color]
+          term_style = STYLES[style]
+
+          # Don't format message if no color or style
+          unless color == :none and style.nil?
+            unless line_prefix.nil?
+              line_prefix = format(line_prefix, term_color, term_style, nil)
+            end
+            message = format(message, term_color, term_style)
+          end
+        end
+
         indent = "%*s" % [MAX_LEVEL, "*" * (MAX_LEVEL - level)]
         (RUBY_VERSION >= "1.9" ? message.lines : message).each do |line|
           if line_prefix
@@ -55,5 +144,10 @@ module Capistrano
     def trace(message, line_prefix=nil)
       log(TRACE, message, line_prefix)
     end
+
+    def format(message, color, style, nl = "\n")
+      style = "#{style};" if style
+      "\e[#{style}#{color}m" + message.to_s.strip + "\e[0m#{nl}"
+    end
   end
 end

test/logger_formatting_test.rb

@@ -0,0 +1,94 @@
+require File.expand_path("../utils", __FILE__)
+require 'capistrano/logger'
+require 'stringio'
+
+Capistrano::Logger.class_eval do
+  # Allows formatters to be changed during tests
+  def self.formatters=(formatters)
+    @formatters = formatters
+    @sorted_formatters = nil
+  end
+end
+
+class LoggerFormattingTest < Test::Unit::TestCase
+  def setup
+    @io = StringIO.new
+    @io.stubs(:tty?).returns(true)
+    @logger = Capistrano::Logger.new(:output => @io, :level => 3)
+  end
+
+  def test_matching_with_style_and_color
+    Capistrano::Logger.formatters = [{ :match => /^err ::/, :color => :red, :style => :underscore, :level => 0 }]
+    @logger.log(0, "err :: Error Occurred")
+    assert @io.string.include? "\e[4;31merr :: Error Occurred\e[0m"
+  end
+
+  def test_style_without_color
+    Capistrano::Logger.formatters = [{ :match => /.*/, :style => :underscore, :level => 0 }]
+    @logger.log(0, "test message")
+    # Default color should be blank (0m)
+    assert @io.string.include? "\e[4;0mtest message\e[0m"
+  end
+
+  def test_prepending_text
+    Capistrano::Logger.formatters = [{ :match => /^executing/, :level => 0, :prepend => '== Currently ' }]
+    @logger.log(0, "executing task")
+    assert @io.string.include? '== Currently executing task'
+  end
+
+  def test_replacing_matched_text
+    Capistrano::Logger.formatters = [{ :match => /^executing/, :level => 0, :replace => 'running' }]
+    @logger.log(0, "executing task")
+    assert @io.string.include? 'running task'
+  end
+
+  def test_prepending_timestamps
+    Capistrano::Logger.formatters = [{ :match => /.*/, :level => 0, :timestamp => true }]
+    @logger.log(0, "test message")
+    assert @io.string.match /\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2} test message/
+  end
+
+  def test_formatter_priorities
+    Capistrano::Logger.formatters = [
+      { :match => /.*/, :color => :red,  :level => 0, :priority => -10 },
+      { :match => /.*/, :color => :blue, :level => 0, :priority => -20, :prepend => '###' }
+    ]
+
+    @logger.log(0, "test message")
+    # Only the red formatter (color 31) should be applied.
+    assert @io.string.include? "\e[31mtest message"
+    # The blue formatter should not have prepended $$$
+    assert !@io.string.include?('###')
+  end
+
+  def test_no_formatting_if_no_color_or_style
+    Capistrano::Logger.formatters = []
+    @logger.log(0, "test message")
+    assert @io.string.include? "*** test message"
+  end
+
+  def test_formatter_log_levels
+    Capistrano::Logger.formatters = [{ :match => /.*/, :color => :blue, :level => 3 }]
+    @logger.log(0, "test message")
+    # Should not match log level
+    assert @io.string.include? "*** test message"
+
+    clear_logger
+    @logger.log(3, "test message")
+    # Should match log level and apply blue color
+    assert @io.string.include? "\e[34mtest message"
+  end
+
+  private
+
+  def colorize(message, color, style = nil)
+    style = "#{style};" if style
+    "\e[#{style}#{color}m" + message + "\e[0m"
+  end
+
+  def clear_logger
+    @io = StringIO.new
+    @io.stubs(:tty?).returns(true)
+    @logger.device = @io
+  end
+end

test/logger_test.rb

@@ -5,7 +5,8 @@ require 'stringio'
 class LoggerTest < Test::Unit::TestCase
   def setup
     @io = StringIO.new
-    @logger = Capistrano::Logger.new(:output => @io)
+    # Turn off formatting for these tests. Formatting is tested in `logger_formatting_test.rb`.
+    @logger = Capistrano::Logger.new(:output => @io, :disable_formatters => true)
   end
 
   def test_logger_should_use_STDERR_by_default