Refactor Colorize#surround

spec/std/colorize_spec.cr

@@ -9,6 +9,14 @@ private def with_color_wrap(*args)
   with_color(*args).toggle(true)
 end
 
+private class ColorizeToS
+  def to_s(io)
+    io << "hello"
+    io << "world".colorize.blue
+    io << "bye"
+  end
+end
+
 describe "colorize" do
   it "colorizes without change" do
     colorize("hello").to_s.should eq("hello")
@@ -121,44 +129,56 @@ describe "colorize" do
     colorize("hello", :red).inspect.should eq("\e[31m\"hello\"\e[0m")
   end
 
-  it "colorizes io with method" do
+  it "colorizes with surround" do
     io = IO::Memory.new
     with_color_wrap.red.surround(io) do
       io << "hello"
+      with_color.green.surround(io) do
+        io << "world"
+      end
+      io << "bye"
     end
-    io.to_s.should eq("\e[31mhello\e[0m")
+    io.to_s.should eq("\e[31mhello\e[0;32mworld\e[0;31mbye\e[0m")
   end
 
-  it "colorizes io with symbol" do
+  it "colorizes with surround and reset" do
     io = IO::Memory.new
-    with_color_wrap(:red).surround(io) do
+    with_color_wrap.red.surround(io) do
       io << "hello"
+      with_color_wrap.green.bold.surround(io) do
+        io << "world"
+      end
+      io << "bye"
     end
-    io.to_s.should eq("\e[31mhello\e[0m")
+    io.to_s.should eq("\e[31mhello\e[0;32;1mworld\e[0;31mbye\e[0m")
   end
 
-  it "colorizes with push and pop" do
+  it "colorizes with surround and no reset" do
     io = IO::Memory.new
-    with_color_wrap.red.push(io) do
+    with_color_wrap.red.surround(io) do
       io << "hello"
-      with_color_wrap.green.push(io) do
+      with_color_wrap.red.surround(io) do
         io << "world"
       end
       io << "bye"
     end
-    io.to_s.should eq("\e[31mhello\e[0;32mworld\e[0;31mbye\e[0m")
+    io.to_s.should eq("\e[31mhelloworldbye\e[0m")
   end
 
-  it "colorizes with push and pop resets" do
+  it "colorizes with surround and default" do
     io = IO::Memory.new
-    with_color_wrap.red.push(io) do
+    with_color_wrap.red.surround(io) do
       io << "hello"
-      with_color_wrap.green.bold.push(io) do
+      with_color_wrap.surround(io) do
         io << "world"
       end
       io << "bye"
     end
-    io.to_s.should eq("\e[31mhello\e[0;32;1mworld\e[0;31mbye\e[0m")
+    io.to_s.should eq("\e[31mhello\e[0mworld\e[31mbye\e[0m")
+  end
+
+  it "colorizes with to_s" do
+    ColorizeToS.new.colorize.red.to_s.should eq("\e[31mhello\e[0;34mworld\e[0;31mbye\e[0m")
   end
 
   it "toggles off" do

src/colorize.cr

@@ -296,64 +296,82 @@ struct Colorize::Object(T)
   end
 
   def surround(io = STDOUT)
-    must_append_end = append_start(io)
-    yield io
-    append_end(io) if must_append_end
-  end
-
-  STACK = [] of Colorize::Object(String)
-
-  def push(io = STDOUT)
-    last_color = STACK.last?
+    return yield io unless @enabled
 
-    append_start(io, !!last_color)
+    Object.surround(io, to_named_tuple) do |io|
+      yield io
+    end
+  end
 
-    STACK.push self
-    yield io
-    STACK.pop
+  private def to_named_tuple
+    {
+      fore: @fore,
+      back: @back,
+      mode: @mode,
+    }
+  end
 
-    if last_color
-      last_color.append_start(io, true)
-    else
-      append_end(io)
+  @@last_color = {
+    fore: ColorANSI::Default.as(Color),
+    back: ColorANSI::Default.as(Color),
+    mode: 0,
+  }
+
+  protected def self.surround(io, color)
+    last_color = @@last_color
+    must_append_end = append_start(io, color)
+    @@last_color = color
+
+    begin
+      yield io
+    ensure
+      append_start(io, last_color) if must_append_end
+      @@last_color = last_color
     end
   end
 
-  protected def append_start(io, reset = false)
-    return false unless @enabled
+  private def self.append_start(io, color)
+    last_color_is_default =
+      @@last_color[:fore] == ColorANSI::Default &&
+        @@last_color[:back] == ColorANSI::Default &&
+        @@last_color[:mode] == 0
+
+    fore = color[:fore]
+    back = color[:back]
+    mode = color[:mode]
 
-    fore_is_default = @fore == ColorANSI::Default
-    back_is_default = @back == ColorANSI::Default
-    mode_is_default = @mode == 0
+    fore_is_default = fore == ColorANSI::Default
+    back_is_default = back == ColorANSI::Default
+    mode_is_default = mode == 0
 
-    if fore_is_default && back_is_default && mode_is_default && !reset
+    if fore_is_default && back_is_default && mode_is_default && last_color_is_default || @@last_color == color
       false
     else
       io << "\e["
 
       printed = false
 
-      if reset
+      unless last_color_is_default
         io << MODE_DEFAULT
         printed = true
       end
 
       unless fore_is_default
         io << ';' if printed
-        @fore.fore io
+        fore.fore io
         printed = true
       end
 
       unless back_is_default
         io << ';' if printed
-        @back.back io
+        back.back io
         printed = true
       end
 
       unless mode_is_default
         # Can't reuse MODES constant because it has bold/bright duplicated
         {% for name in %w(bold dim underline blink reverse hidden) %}
-          if @mode.bits_set? MODE_{{name.upcase.id}}_FLAG
+          if mode.bits_set? MODE_{{name.upcase.id}}_FLAG
             io << ';' if printed
             io << MODE_{{name.upcase.id}}
             printed = true
@@ -366,8 +384,4 @@ struct Colorize::Object(T)
       true
     end
   end
-
-  protected def append_end(io)
-    Colorize.reset(io)
-  end
 end

src/compiler/crystal/tools/print_hierarchy.cr

@@ -27,7 +27,7 @@ module Crystal
         compute_targets(@program.types, exp, false)
       end
 
-      with_color.light_gray.bold.push(STDOUT) do
+      with_color.light_gray.bold.surround(STDOUT) do
         print_type @program.object
       end
     end
@@ -126,7 +126,7 @@ module Crystal
       if (type.is_a?(NonGenericClassType) || type.is_a?(GenericClassInstanceType)) &&
          !type.is_a?(PointerInstanceType) && !type.is_a?(ProcInstanceType)
         size = @llvm_typer.size_of(@llvm_typer.llvm_struct_type(type))
-        with_color.light_gray.push(STDOUT) do
+        with_color.light_gray.surround(STDOUT) do
           print " ("
           print size.to_s
           print " bytes)"
@@ -176,7 +176,7 @@ module Crystal
           print "      "
         end
 
-        with_color.light_gray.push(STDOUT) do
+        with_color.light_gray.surround(STDOUT) do
           print name.ljust(max_name_size)
           print " : "
           print var
@@ -211,13 +211,13 @@ module Crystal
           print "      "
         end
 
-        with_color.light_gray.push(STDOUT) do
+        with_color.light_gray.surround(STDOUT) do
           print ivar.name.ljust(max_name_size)
           print " : "
           if ivar_type = ivar.type?
             print ivar_type.to_s.ljust(max_type_size)
             size = @llvm_typer.size_of(@llvm_typer.llvm_embedded_type(ivar_type))
-            with_color.light_gray.push(STDOUT) do
+            with_color.light_gray.surround(STDOUT) do
               print " ("
               print size.to_s.rjust(max_bytes_size)
               print " bytes)"