From 7907a5871d675cec2294f159d537360b82e2e6c6 Mon Sep 17 00:00:00 2001
From: Andrew O'Brien <andrew@econify.com>
Date: Wed, 4 May 2011 20:20:07 -0400
Subject: [PATCH] Added AwesomePrint.force_colors! to allow outputting colors
 when not run in a process attached to a TTY.

---
 lib/ap/awesome_print.rb   | 11 +++++-
 lib/ap/core_ext/string.rb | 20 +++++++---
 spec/colorization_spec.rb | 80 +++++++++++++++++++++++++++++++++++++++
 spec/spec_helper.rb       |  4 ++
 4 files changed, 108 insertions(+), 7 deletions(-)
 create mode 100644 spec/colorization_spec.rb

diff --git a/lib/ap/awesome_print.rb b/lib/ap/awesome_print.rb
index e17fb36..6dcf0e8 100755
--- a/lib/ap/awesome_print.rb
+++ b/lib/ap/awesome_print.rb
@@ -43,7 +43,16 @@ class AwesomePrint
     @indentation = @options[:indent].abs
     Thread.current[AP] ||= []
   end
-  
+
+  def self.allow_colors?
+    (STDOUT.tty? && ((ENV['TERM'] && ENV['TERM'] != 'dumb') || ENV['ANSICON'])) || @@force_colors
+  end
+
+  @@force_colors = false
+  def self.force_colors!(value = true)
+    @@force_colors = value
+  end
+
   private
 
   # Format an array.
diff --git a/lib/ap/core_ext/string.rb b/lib/ap/core_ext/string.rb
index 60fe954..8dbede1 100644
--- a/lib/ap/core_ext/string.rb
+++ b/lib/ap/core_ext/string.rb
@@ -12,12 +12,20 @@ class String
   #    0 => normal
 
   [ :gray, :red, :green, :yellow, :blue, :purple, :cyan, :white ].each_with_index do |color, i|
-    if STDOUT.tty? && ((ENV['TERM'] && ENV['TERM'] != 'dumb') || ENV['ANSICON'])
-      define_method color          do "\033[1;#{30+i}m#{self}\033[0m" end
-      define_method :"#{color}ish" do "\033[0;#{30+i}m#{self}\033[0m" end
-    else
-      define_method color do self end
-      alias_method :"#{color}ish", color 
+    define_method color do
+      if AwesomePrint.allow_colors?
+        "\033[1;#{30+i}m#{self}\033[0m"
+      else
+        self
+      end
+    end
+
+    define_method :"#{color}ish" do
+      if AwesomePrint.allow_colors?
+        "\033[0;#{30+i}m#{self}\033[0m"
+      else
+        self
+      end
     end
   end
 
diff --git a/spec/colorization_spec.rb b/spec/colorization_spec.rb
new file mode 100644
index 0000000..be96869
--- /dev/null
+++ b/spec/colorization_spec.rb
@@ -0,0 +1,80 @@
+require File.expand_path(File.dirname(__FILE__) + '/spec_helper')
+
+describe "AwesomePrint" do
+  describe "colorization" do
+    PLAIN = '[ 1, :two, "three", [ nil, [ true, false ] ] ]'
+    COLORIZED = "[ \e[1;34m1\e[0m, \e[0;36m:two\e[0m, \e[0;33m\"three\"\e[0m, [ \e[1;31mnil\e[0m, [ \e[1;32mtrue\e[0m, \e[1;31mfalse\e[0m ] ] ]"
+
+    before(:each) do
+      AwesomePrint.force_colors!(false)
+      ENV['TERM'] = "xterm-colors"
+      ENV.delete('ANSICON')
+      @arr = [ 1, :two, "three", [ nil, [ true, false] ] ]
+    end
+    
+    it "colorizes tty processes by default" do
+      stub_tty!(STDOUT, true)
+
+      @arr.ai(:multiline => false).should == COLORIZED
+    end
+
+    it "colorizes tty processes by default" do
+      stub_tty!(STDOUT, true)
+
+      @arr.ai(:multiline => false).should == COLORIZED
+    end
+
+    
+    it "colorizes processes with ENV['ANSICON'] by default" do
+      stub_tty!(STDOUT, true)
+      ENV['ANSICON'] = "1"
+
+      @arr.ai(:multiline => false).should == COLORIZED
+    end
+
+    it "does not colorize tty processes running in dumb terminals by default" do
+      stub_tty!(STDOUT, true)
+      ENV['TERM'] = "dumb"
+
+      @arr.ai(:multiline => false).should == PLAIN
+    end
+
+    it "does not colorize subprocesses by default" do
+      stub_tty!(STDOUT, false)
+
+      @arr.ai(:multiline => false).should == PLAIN
+    end
+    
+    describe "forced" do
+      before(:each) do
+        AwesomePrint.force_colors!
+      end
+      
+      it "still colorizes tty processes" do
+        stub_tty!(STDOUT, true)
+
+        @arr.ai(:multiline => false).should == COLORIZED
+      end
+      
+      it "colorizes dumb terminals" do
+        stub_tty!(STDOUT, true)
+        ENV["TERM"] = "dumb"
+
+        @arr.ai(:multiline => false).should == COLORIZED
+      end
+
+      it "colorizes subprocess" do
+        stub_tty!(STDOUT, true)
+        @arr.ai(:multiline => false).should == COLORIZED
+      end
+    end
+  end
+  
+  def stub_tty!(stream, value)
+    eval(%{class << stream
+      def tty?
+        #{value}
+      end
+    end})
+  end
+end
\ No newline at end of file
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 3363ea7..b518b8b 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -26,3 +26,7 @@ def stub_dotfile!
   dotfile = File.join(ENV["HOME"], ".aprc")
   File.should_receive(:readable?).at_least(:once).with(dotfile).and_return(false)
 end
+
+# Infinity Test runs tests as subprocesses, which sets STDOUT.tty? to false and 
+# would otherwise prematurely disallow colors. We'll test the defaults later.
+AwesomePrint.force_colors!