Merge pull request #46 from red-data-tools/render_color_param

Let `render` method support `color:` option

Author: mrkn

lib/unicode_plot.rb

@@ -2,6 +2,7 @@
 
 require 'unicode_plot/version'
 
+require 'unicode_plot/io_context'
 require 'unicode_plot/utils'
 require 'unicode_plot/styled_printer'
 require 'unicode_plot/value_transformer'

lib/unicode_plot/io_context.rb

@@ -0,0 +1,32 @@
+require "forwardable"
+
+module UnicodePlot
+  class IOContext
+    extend Forwardable
+
+    def initialize(io, color: :auto)
+      @io = io
+      @color = check_color(color)
+    end
+
+    def_delegators :@io, :print, :puts
+
+    def color?
+      case @color
+      when :auto
+        @io.respond_to?(:tty?) ? @io.tty? : false
+      else
+        @color
+      end
+    end
+
+    private def check_color(color)
+      case color
+      when true, false, :auto
+        color
+      else
+        raise ArgumentError, "color must be either true, false, :auto"
+      end
+    end
+  end
+end

lib/unicode_plot/plot.rb

@@ -103,8 +103,8 @@ def annotate_row!(loc, row_index, value, color: :normal)
       end
     end
 
-    def render(out=$stdout, newline: true)
-      Renderer.render(out, self, newline)
+    def render(out=$stdout, newline: true, color: :auto)
+      Renderer.render(IOContext.new(out, color: color), self, newline)
     end
 
     COLOR_CYCLE = [

lib/unicode_plot/renderer.rb

@@ -128,7 +128,7 @@ def render_row(row)
       left_len  = nocolor_string(left_str).length
       right_len = nocolor_string(right_str).length
 
-      unless color?(out)
+      unless out.color?
         left_str  = nocolor_string(left_str)
         right_str = nocolor_string(right_str)
       end

lib/unicode_plot/styled_printer.rb

@@ -58,7 +58,7 @@ module StyledPrinter
     COLOR_DECODE = COLOR_ENCODE.map {|k, v| [v, k] }.to_h.freeze
 
     def print_styled(out, *args, bold: false, color: :normal)
-      return out.print(*args) unless color?(out)
+      return out.print(*args) unless out.color?
 
       str = StringIO.open {|sio| sio.print(*args); sio.close; sio.string }
       color = :nothing if bold && color == :bold
@@ -83,9 +83,5 @@ def print_color(out, color, *args)
       color = COLOR_DECODE[color]
       print_styled(out, *args, color: color)
     end
-
-    def color?(out)
-      out&.tty? || false
-    end
   end
 end

test/helper/with_term.rb

@@ -2,22 +2,28 @@
 
 module Helper
   module WithTerm
-    def with_term(tty=true)
+    def with_sio(tty: true)
       sio = StringIO.new
       def sio.tty?; true; end if tty
 
-      orig_stdout, $stdout = $stdout, sio
-      orig_env = ENV.to_h.dup
-      ENV['TERM'] = 'xterm-256color'
-
-      result = yield
+      result = yield(sio)
       sio.close
 
       [result, sio.string]
-    ensure
-      $stdout.close
-      $stdout = orig_stdout
-      ENV.replace(orig_env) if orig_env
+    end
+
+    def with_term(tty=true)
+      with_sio(tty: tty) do |sio|
+        begin
+          orig_stdout, $stdout = $stdout, sio
+          orig_env = ENV.to_h.dup
+          ENV['TERM'] = 'xterm-256color'
+          yield
+        ensure
+          $stdout = orig_stdout
+          ENV.replace(orig_env) if orig_env
+        end
+      end
     end
   end
 end

test/test-canvas.rb

@@ -55,11 +55,11 @@ def setup
 
       test("empty") do
         if self.class::CANVAS_NAME == :braille
-          _, output = with_term { @canvas.show($stdout) }
+          _, output = with_term { @canvas.show(UnicodePlot::IOContext.new($stdout)) }
           assert_equal(fixture_path("canvas/empty_braille_show.txt").read,
                        output)
         else
-          _, output = with_term { @canvas.show($stdout) }
+          _, output = with_term { @canvas.show(UnicodePlot::IOContext.new($stdout)) }
           assert_equal(fixture_path("canvas/empty_show.txt").read,
                        output)
         end
@@ -81,31 +81,31 @@ def setup
         end
 
         test("print_row") do
-          _, output = with_term { @canvas.print_row($stdout, 2) }
+          _, output = with_term { @canvas.print_row(UnicodePlot::IOContext.new($stdout), 2) }
           assert_equal(fixture_path("canvas/#{self.class::CANVAS_NAME}_printrow.txt").read,
                        output)
         end
 
         test("print") do
-          _, output = with_term { @canvas.print($stdout) }
+          _, output = with_term { @canvas.print(UnicodePlot::IOContext.new($stdout)) }
           assert_equal(fixture_path("canvas/#{self.class::CANVAS_NAME}_print.txt").read,
                        output)
         end
 
         test("print_nocolor") do
-          _, output = with_term(false) { @canvas.print($stdout) }
+          _, output = with_term(false) { @canvas.print(UnicodePlot::IOContext.new($stdout)) }
           assert_equal(fixture_path("canvas/#{self.class::CANVAS_NAME}_print_nocolor.txt").read,
                        output)
         end
 
         test("sow") do
-          _, output = with_term { @canvas.show($stdout) }
+          _, output = with_term { @canvas.show(UnicodePlot::IOContext.new($stdout)) }
           assert_equal(fixture_path("canvas/#{self.class::CANVAS_NAME}_show.txt").read,
                        output)
         end
 
         test("show_nocolor") do
-          _, output = with_term(false) { @canvas.show($stdout) }
+          _, output = with_term(false) { @canvas.show(UnicodePlot::IOContext.new($stdout)) }
           assert_equal(fixture_path("canvas/#{self.class::CANVAS_NAME}_show_nocolor.txt").read,
                        output)
         end

test/test-plot.rb

@@ -1,6 +1,8 @@
 require 'stringio'
 
 class TestPlot < Test::Unit::TestCase
+  include Helper::WithTerm
+
   sub_test_case("#render") do
     test("render to $stdout when no arguments") do
       sio = StringIO.new
@@ -16,5 +18,19 @@ class TestPlot < Test::Unit::TestCase
         $stdout = save_stdout
       end
     end
+
+    test("color: true") do
+      _, tty_output   = with_sio(tty: true)  {|sio| UnicodePlot.barplot(data: {a: 23, b: 37}).render(sio) }
+      _, notty_output = with_sio(tty: false) {|sio| UnicodePlot.barplot(data: {a: 23, b: 37}).render(sio, color: true) }
+
+      assert_equal(tty_output, notty_output)
+    end
+
+    test("color: false") do
+      _, tty_output   = with_sio(tty: true)  {|sio| UnicodePlot.barplot(data: {a: 23, b: 37}).render(sio, color: false) }
+      _, notty_output = with_sio(tty: false) {|sio| UnicodePlot.barplot(data: {a: 23, b: 37}).render(sio) }
+
+      assert_equal(tty_output, notty_output)
+    end
   end
 end