lib/strings/wrap.rb in strings-0.1.5 vs lib/strings/wrap.rb in strings-0.1.6

@@ -7,35 +7,34 @@
 require_relative 'fold'
 
 module Strings
   module Wrap
     DEFAULT_WIDTH = 80
-
     NEWLINE = "\n".freeze
+    SPACE = " ".freeze
+    LINE_BREAK = %r{\r\n|\r|\n}.freeze
+    LINE_BREAKS = "\r\n+|\r+|\n+".freeze
 
-    SPACE = ' '.freeze
-
-    LINE_BREAK = "\r\n+|\r+|\n+".freeze
-
     # Wrap a text into lines no longer than wrap_at length.
     # Preserves existing lines and existing word boundaries.
     #
     # @example
     #   Strings::Wrap.wrap("Some longish text", 8)
     #   # => >Some
     #        >longish
     #        >text
     #
     # @api public
-    def wrap(text, wrap_at = DEFAULT_WIDTH)
-      if text.length < wrap_at.to_i || wrap_at.to_i.zero?
+    def wrap(text, wrap_at = DEFAULT_WIDTH, separator: nil)
+      if text.scan(/[[:print:]]/).length < wrap_at.to_i || wrap_at.to_i.zero?
         return text
       end
       ansi_stack = []
-      text.split(%r{#{LINE_BREAK}}, -1).map do |paragraph|
+      sep = separator || text[LINE_BREAK] || NEWLINE
+      text.split(%r{#{LINE_BREAKS}}, -1).map do |paragraph|
         format_paragraph(paragraph, wrap_at, ansi_stack)
-      end * NEWLINE
+      end * sep
     end
     module_function :wrap
 
     # Format paragraph to be maximum of wrap_at length
     #
@@ -68,11 +67,16 @@
           if Strings::ANSI.only_ansi?(ansi.join)
             ansi_matched = true
           elsif ansi_matched
             ansi_stack << [ansi[0...-1].join, line_length + word_length]
             ansi_matched = false
-            ansi = []
+
+            if ansi.last == Strings::ANSI::CSI
+              ansi = [ansi.last]
+            else
+              ansi = []
+            end
           end
           next if ansi.length > 0
         end
 
         char_length = display_width(char)
@@ -89,68 +93,75 @@
           end
           next
         end
 
         if char == SPACE # ends with space
-          lines << insert_ansi(ansi_stack, line.join)
+          lines << insert_ansi(line.join, ansi_stack)
           line = []
           line_length = 0
           word << char
           word_length += char_length
         elsif word_length + char_length <= wrap_at
-          lines << insert_ansi(ansi_stack, line.join)
+          lines << insert_ansi(line.join, ansi_stack)
           line = [word.join + char]
           line_length = word_length + char_length
           word = []
           word_length = 0
         else # hyphenate word - too long to fit a line
-          lines << insert_ansi(ansi_stack, word.join)
+          lines << insert_ansi(word.join, ansi_stack)
           line_length = 0
           word = [char]
           word_length = char_length
         end
       end
-      lines << insert_ansi(ansi_stack, line.join) unless line.empty?
-      lines << insert_ansi(ansi_stack, word.join) unless word.empty?
+      lines << insert_ansi(line.join, ansi_stack) unless line.empty?
+      lines << insert_ansi(word.join, ansi_stack) unless word.empty?
       lines
     end
     module_function :format_paragraph
 
     # Insert ANSI code into string
     #
     # Check if there are any ANSI states, if present
     # insert ANSI codes at given positions unwinding the stack.
     #
-    # @param [Array[Array[String, Integer]]] ansi_stack
-    #   the ANSI codes to apply
-    #
     # @param [String] string
     #   the string to insert ANSI codes into
     #
+    # @param [Array[Array[String, Integer]]] ansi_stack
+    #   the ANSI codes to apply
+    #
     # @return [String]
     #
     # @api private
-    def insert_ansi(ansi_stack, string)
+    def insert_ansi(string, ansi_stack = [])
       return string if ansi_stack.empty?
-      to_remove = 0
-      reset_index = -1
-      output = string.dup
-      resetting = false
-      ansi_stack.reverse_each do |state|
-        if state[0] =~ /#{Regexp.quote(Strings::ANSI::RESET)}/
-          resetting = true
-          reset_index = state[1]
-          to_remove += 2
+      return string if string.empty?
+
+      new_stack = []
+      output          = string.dup
+      length          = string.size
+      matched_reset   = false
+      ansi_reset      = Strings::ANSI::RESET
+
+      # Reversed so that string index don't count ansi
+      ansi_stack.reverse_each do |ansi|
+        if ansi[0] =~ /#{Regexp.quote(ansi_reset)}/
+          matched_reset = true
+          output.insert(ansi[1], ansi_reset)
           next
-        elsif !resetting
-          reset_index = -1
-          resetting = false
+        elsif !matched_reset # ansi without reset
+          matched_reset = false
+          new_stack << ansi # keep the ansi
+          next if ansi[1] == length
+          output.insert(-1, ansi_reset) # add reset at the end
         end
 
-        color, color_index = *state
-        output.insert(reset_index, Strings::ANSI::RESET).insert(color_index, color)
+        output.insert(ansi[1], ansi[0])
       end
-      ansi_stack.pop(to_remove) # remove used states
+
+      ansi_stack.replace(new_stack)
+
       output
     end
     module_function :insert_ansi
 
     # Visible width of a string