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/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

@@ -16,5 +16,16 @@ class TestPlot < Test::Unit::TestCase
         $stdout = save_stdout
       end
     end
+
+    test("color: true") do
+      sio_notty = StringIO.new
+      UnicodePlot.barplot(data: {a: 23, b: 37}).render(sio_notty, color: true)
+
+      class << (sio_tty = StringIO.new)
+        def tty?; true; end
+      end
+      UnicodePlot.barplot(data: {a: 23, b: 37}).render(sio_tty)
+      assert_equal(sio_tty.string, sio_notty.string)
+    end
   end
 end