Experimental: refactor Airbrussh::Colors to be class, rather than module

lib/airbrussh/capistrano/tasks.rb

@@ -16,13 +16,13 @@ class Tasks
       extend Forwardable
       def_delegators :dsl, :set
       def_delegators :config, :log_file
-
-      include Airbrussh::Colors
+      def_delegators :@colors, *Airbrussh::Colors.names
 
       def initialize(dsl, stderr=$stderr, config=Airbrussh.configuration)
         @dsl = dsl
         @stderr = stderr
         @config = config
+        @colors = config.colors(stderr)
 
         configure
         warn_if_missing_dsl

lib/airbrussh/colors.rb

@@ -1,23 +1,47 @@
 module Airbrussh
   # Very basic support for ANSI color, so that we don't have to rely on
   # any external dependencies.
-  module Colors
+  class Colors
     ANSI_CODES = {
-      :red    => 31,
-      :green  => 32,
-      :yellow => 33,
-      :blue   => 34,
-      :gray   => 90
+      :black   => 30,
+      :red     => 31,
+      :green   => 32,
+      :yellow  => 33,
+      :blue    => 34,
+      :magenta => 35,
+      :cyan    => 36,
+      :white   => 37,
+      :gray    => 90,
+      :light_black   => 90,
+      :light_red     => 91,
+      :light_green   => 92,
+      :light_yellow  => 93,
+      :light_blue    => 94,
+      :light_magenta => 95,
+      :light_cyan    => 96,
+      :light_white   => 97
     }.freeze
 
+    def self.names
+      ANSI_CODES.keys
+    end
+
+    def initialize(enabled=true)
+      @enabled = enabled
+    end
+
+    def enabled?
+      @enabled
+    end
+
     # Define red, green, blue, etc. methods that return a copy of the
     # String that is wrapped in the corresponding ANSI color escape
     # sequence.
     ANSI_CODES.each do |name, code|
       define_method(name) do |string|
+        return string.to_s unless enabled?
         "\e[0;#{code};49m#{string}\e[0m"
       end
-      module_function(name)
     end
   end
 end

lib/airbrussh/command_formatter.rb

@@ -1,6 +1,6 @@
 # encoding: UTF-8
 require "airbrussh/colors"
-require "delegate"
+require "forwardable"
 # rubocop:disable Style/AsciiComments
 
 module Airbrussh
@@ -11,11 +11,13 @@ module Airbrussh
   #              all commands that have been run in the current rake task
   #  class CommandFormatter < SimpleDelegator
   class CommandFormatter < SimpleDelegator
-    include Airbrussh::Colors
+    extend Forwardable
+    def_delegators :@colors, *Airbrussh::Colors.names
 
-    def initialize(command, position)
+    def initialize(command, position, colors=Airbrussh::Colors.new)
       super(command)
       @position = position
+      @colors = colors
     end
 
     # Prefixes the line with the command number and removes the newline.

lib/airbrussh/configuration.rb

@@ -16,15 +16,22 @@ def initialize
       self.command_output = false
     end
 
-    def banner_message
+    # Returns a configured Colors object appropriate for the given io.
+    # This is based on the color setting (true, false, or :auto), and whether
+    # the io is a tty.
+    def colors(io)
+      Airbrussh::Colors.new(color?(io))
+    end
+
+    def banner_message(io)
       return nil unless banner
       return banner unless banner == :auto
-      msg = "Using airbrussh format."
+      m = "Using airbrussh format."
       if log_file
-        msg << "\n"
-        msg << "Verbose output is being written to #{Colors.blue(log_file)}."
+        m << "\n"
+        m << "Verbose output is being written to #{colors(io).blue(log_file)}."
       end
-      msg
+      m
     end
 
     # This returns an array of formatters appropriate for the configuration.
@@ -39,5 +46,18 @@ def formatters(io)
     def show_command_output?(sym)
       command_output == true || Array(command_output).include?(sym)
     end
+
+    private
+
+    def color?(io)
+      case color
+      when true
+        true
+      when :auto
+        ENV["SSHKIT_COLOR"] || (io.respond_to?("tty?") && io.tty?)
+      else
+        false
+      end
+    end
   end
 end

lib/airbrussh/console.rb

@@ -5,9 +5,9 @@ module Airbrussh
   # Helper class that wraps an IO object and provides methods for truncating
   # output, assuming the IO object represents a console window.
   #
-  # This is useful for writing log messages that will typically show up on
-  # an ANSI color-capable console. When a console is not present (e.g. when
-  # running on a CI server) the output will gracefully degrade.
+  # If color is disabled for the IO object (based on Airbrussh::Configuration),
+  # any ANSI color codes will also be stripped from the output.
+  #
   class Console
     attr_reader :output, :config
 
@@ -17,9 +17,8 @@ def initialize(output, config=Airbrussh.configuration)
     end
 
     # Writes to the IO after first truncating the output to fit the console
-    # width. If the underlying IO is not a TTY, ANSI colors are removed from
-    # the output. A newline is always added. Color output can be forced by
-    # setting the SSHKIT_COLOR environment variable.
+    # width. If color is disabled for the underlying IO, ANSI colors are also
+    # removed from the output. A newline is always added.
     def print_line(obj="")
       string = obj.to_s
 
@@ -67,14 +66,7 @@ def console_width
     private
 
     def color_enabled?
-      case config.color
-      when true
-        true
-      when :auto
-        ENV["SSHKIT_COLOR"] || @output.tty?
-      else
-        false
-      end
+      @color_enabled ||= config.colors(output).enabled?
     end
 
     def utf8_supported?(string)

lib/airbrussh/console_formatter.rb

@@ -2,12 +2,13 @@
 require "airbrussh/command_formatter"
 require "airbrussh/console"
 require "airbrussh/rake/context"
+require "forwardable"
 require "sshkit"
 
 module Airbrussh
   class ConsoleFormatter < SSHKit::Formatter::Abstract
-    include Airbrussh::Colors
     extend Forwardable
+    def_delegators :@colors, *Airbrussh::Colors.names
 
     attr_reader :config, :context
     def_delegators :context, :current_task_name, :register_new_command
@@ -16,14 +17,16 @@ def initialize(io, config=Airbrussh.configuration)
       super(io)
 
       @config = config
+      @colors = config.colors(io)
       @context = Airbrussh::Rake::Context.new(config)
       @console = Airbrussh::Console.new(original_output, config)
 
       write_banner
     end
 
     def write_banner
-      print_line(config.banner_message) if config.banner_message
+      message = config.banner_message(original_output)
+      print_line(message) if message
     end
 
     def log_command_start(command)
@@ -105,7 +108,11 @@ def debug?(obj)
     end
 
     def decorate(command)
-      Airbrussh::CommandFormatter.new(command, @context.position(command))
+      Airbrussh::CommandFormatter.new(
+        command,
+        @context.position(command),
+        @colors
+      )
     end
 
     def print_line(string)

test/airbrussh/colors_test.rb

@@ -1,26 +1,64 @@
 require "minitest_helper"
 require "airbrussh/colors"
+require "forwardable"
 
 class Airbrussh::ColorsTest < Minitest::Test
-  include Airbrussh::Colors
+  extend Forwardable
+  def_delegators :@colors, *Airbrussh::Colors.names
 
   def test_red
+    enable_color
     assert_equal("\e[0;31;49mhello\e[0m", red("hello"))
   end
 
   def test_green
+    enable_color
     assert_equal("\e[0;32;49mhello\e[0m", green("hello"))
   end
 
   def test_yellow
+    enable_color
     assert_equal("\e[0;33;49mhello\e[0m", yellow("hello"))
   end
 
   def test_blue
+    enable_color
     assert_equal("\e[0;34;49mhello\e[0m", blue("hello"))
   end
 
   def test_gray
+    enable_color
     assert_equal("\e[0;90;49mhello\e[0m", gray("hello"))
   end
+
+  def test_red_disabled
+    enable_color(false)
+    assert_equal("hello", red("hello"))
+  end
+
+  def test_green_disabled
+    enable_color(false)
+    assert_equal("hello", green("hello"))
+  end
+
+  def test_yellow_disabled
+    enable_color(false)
+    assert_equal("hello", yellow("hello"))
+  end
+
+  def test_blue_disabled
+    enable_color(false)
+    assert_equal("hello", blue("hello"))
+  end
+
+  def test_gray_disabled
+    enable_color(false)
+    assert_equal("hello", gray("hello"))
+  end
+
+  private
+
+  def enable_color(enabled=true)
+    @colors = Airbrussh::Colors.new(enabled)
+  end
 end

test/airbrussh/configuration_test.rb

@@ -20,29 +20,30 @@ def test_defaults
   def test_auto_banner_message_without_log
     @config.log_file = nil
     @config.banner = :auto
-    assert_equal("Using airbrussh format.", @config.banner_message)
+    assert_equal("Using airbrussh format.", @config.banner_message(stub))
   end
 
   def test_auto_banner_message_with_log
     @config.log_file = "log/test.log"
     @config.banner = :auto
+    @config.color = true
     assert_equal(
       "Using airbrussh format.\n"\
       "Verbose output is being written to \e[0;34;49mlog/test.log\e[0m.",
-      @config.banner_message
+      @config.banner_message(stub)
     )
   end
 
   def test_nil_or_false_banner_message
     @config.banner = nil
-    assert_nil(@config.banner_message)
+    assert_nil(@config.banner_message(stub))
     @config.banner = false
-    assert_nil(@config.banner_message)
+    assert_nil(@config.banner_message(stub))
   end
 
   def test_custom_banner_message
     @config.banner = "Hello!"
-    assert_equal("Hello!", @config.banner_message)
+    assert_equal("Hello!", @config.banner_message(stub))
   end
 
   def test_formatters_without_log_file
@@ -95,4 +96,30 @@ def test_effects_of_command_output_stdout_stderr
     assert(@config.show_command_output?(:stdout))
     assert(@config.show_command_output?(:stderr))
   end
+
+  def test_color_is_allowed_for_tty
+    @config.color = :auto
+    assert(@config.colors(stub("tty?" => true)).enabled?)
+  end
+
+  def test_color_can_be_forced
+    @config.color = true
+    assert(@config.colors(stub("tty?" => false)).enabled?)
+  end
+
+  def test_color_can_be_forced_via_env
+    @config.color = :auto
+    ENV.stubs(:[]).with("SSHKIT_COLOR").returns("1")
+    assert(@config.colors(stub("tty?" => false)).enabled?)
+  end
+
+  def test_color_disabled_non_tty
+    @config.color = :auto
+    refute(@config.colors(stub("tty?" => false)).enabled?)
+  end
+
+  def test_color_can_be_disabled_for_tty
+    @config.color = false
+    refute(@config.colors(stub("tty?" => true)).enabled?)
+  end
 end