Merge pull request #19 from cprice-puppet/feature/master/use-existing-indentation

Feature/master/use existing indentation

2 files changed, 85 insertions, 14 deletions

diff --git a/lib/puppet/util/ini_file.rb b/lib/puppet/util/ini_file.rb
index 52ad32c..0409db2 100644
--- a/lib/puppet/util/ini_file.rb
+++ b/lib/puppet/util/ini_file.rb
@@ -6,7 +6,7 @@ module Util
   class IniFile
 
     SECTION_REGEX = /^\s*\[([\w\d\.\\\/\-\:]+)\]\s*$/
-    SETTING_REGEX = /^\s*([\w\d\.\\\/\-]+)\s*=\s*([\S\s]*\S)\s*$/
+    SETTING_REGEX = /^(\s*)([\w\d\.\\\/\-]+)(\s*=\s*)([\S\s]*\S)\s*$/
 
     def initialize(path, key_val_separator = ' = ')
       @path = path
@@ -30,7 +30,7 @@ module Util
 
     def set_value(section_name, setting, value)
       unless (@sections_hash.has_key?(section_name))
-        add_section(Section.new(section_name, nil, nil, nil))
+        add_section(Section.new(section_name, nil, nil, nil, nil))
       end
 
       section = @sections_hash[section_name]
@@ -64,21 +64,62 @@ module Util
     def save
       File.open(@path, 'w') do |fh|
 
-        @section_names.each do |name|
+        @section_names.each_index do |index|
+          name = @section_names[index]
 
           section = @sections_hash[name]
 
-          if section.start_line.nil?
+          # We need a buffer to cache lines that are only whitespace
+          whitespace_buffer = []
+
+          if (section.is_new_section?) && (! section.is_global?)
             fh.puts("\n[#{section.name}]")
-          elsif ! section.end_line.nil?
+          end
+
+          if ! section.is_new_section?
+            # write all of the pre-existing settings
             (section.start_line..section.end_line).each do |line_num|
-              fh.puts(lines[line_num])
+              line = lines[line_num]
+
+              # We buffer any lines that are only whitespace so that
+              # if they are at the end of a section, we can insert
+              # any new settings *before* the final chunk of whitespace
+              # lines.
+              if (line =~ /^\s*$/)
+                whitespace_buffer << line
+              else
+                # If we get here, we've found a non-whitespace line.
+                # We'll flush any cached whitespace lines before we
+                # write it.
+                flush_buffer_to_file(whitespace_buffer, fh)
+                fh.puts(line)
+              end
             end
           end
 
+          # write new settings, if there are any
           section.additional_settings.each_pair do |key, value|
-            fh.puts("#{key}#{@key_val_separator}#{value}")
+            fh.puts("#{' ' * (section.indentation || 0)}#{key}#{@key_val_separator}#{value}")
+          end
+
+          if (whitespace_buffer.length > 0)
+            flush_buffer_to_file(whitespace_buffer, fh)
+          else
+            # We get here if there were no blank lines at the end of the
+            # section.
+            #
+            # If we are adding a new section with a new setting,
+            # and if there are more sections that come after this one,
+            # we'll write one blank line just so that there is a little
+            # whitespace between the sections.
+            #if (section.end_line.nil? &&
+            if (section.is_new_section? &&
+                (section.additional_settings.length > 0) &&
+                (index < @section_names.length - 1))
+              fh.puts("")
+            end
           end
+
         end
       end
     end
@@ -111,12 +152,15 @@ module Util
     def read_section(name, start_line, line_iter)
       settings = {}
       end_line_num = nil
+      min_indentation = nil
       while true
         line, line_num = line_iter.peek
         if (line_num.nil? or match = SECTION_REGEX.match(line))
-          return Section.new(name, start_line, end_line_num, settings)
+          return Section.new(name, start_line, end_line_num, settings, min_indentation)
         elsif (match = SETTING_REGEX.match(line))
-          settings[match[1]] = match[2]
+          settings[match[2]] = match[4]
+          indentation = match[1].length
+          min_indentation = [indentation, min_indentation || indentation].min
         end
         end_line_num = line_num
         line_iter.next
@@ -126,8 +170,8 @@ module Util
     def update_line(section, setting, value)
       (section.start_line..section.end_line).each do |line_num|
         if (match = SETTING_REGEX.match(lines[line_num]))
-          if (match[1] == setting)
-            lines[line_num] = "#{setting}#{@key_val_separator}#{value}"
+          if (match[2] == setting)
+            lines[line_num] = "#{match[1]}#{match[2]}#{match[3]}#{value}"
           end
         end
       end
@@ -136,7 +180,7 @@ module Util
     def remove_line(section, setting)
       (section.start_line..section.end_line).each do |line_num|
         if (match = SETTING_REGEX.match(lines[line_num]))
-          if (match[1] == setting)
+          if (match[2] == setting)
             lines.delete_at(line_num)
           end
         end
@@ -173,6 +217,13 @@ module Util
       end
     end
 
+    def flush_buffer_to_file(buffer, fh)
+      if buffer.length > 0
+        buffer.each { |l| fh.puts(l) }
+        buffer.clear
+      end
+    end
+
   end
 end
 end
diff --git a/lib/puppet/util/ini_file/section.rb b/lib/puppet/util/ini_file/section.rb
index 16f19d3..d7ff159 100644
--- a/lib/puppet/util/ini_file/section.rb
+++ b/lib/puppet/util/ini_file/section.rb
@@ -2,15 +2,35 @@ module Puppet
 module Util
 class IniFile
   class Section
-    def initialize(name, start_line, end_line, settings)
+    # Some implementation details:
+    #
+    #  * `name` will be set to the empty string for the 'global' section.
+    #  * there will always be a 'global' section, with a `start_line` of 0,
+    #    but if the file actually begins with a real section header on
+    #    the first line, then the 'global' section will have an
+    #    `end_line` of `nil`.
+    #  * `start_line` and `end_line` will be set to `nil` for a new non-global
+    #    section.
+    def initialize(name, start_line, end_line, settings, indentation)
       @name = name
       @start_line = start_line
       @end_line = end_line
       @existing_settings = settings.nil? ? {} : settings
       @additional_settings = {}
+      @indentation = indentation
     end
 
-    attr_reader :name, :start_line, :end_line, :additional_settings
+    attr_reader :name, :start_line, :end_line, :additional_settings, :indentation
+
+    def is_global?()
+      @name == ''
+    end
+
+    def is_new_section?()
+      # a new section (global or named) will always have `end_line`
+      # set to `nil`
+      @end_line.nil?
+    end
 
     def get_value(setting_name)
       @existing_settings[setting_name] || @additional_settings[setting_name]