jm · pinecat · Feb 13, 2022
diff --git a/lib/toml/generator.rb b/lib/toml/generator.rb
@@ -1,19 +1,18 @@
-
-module TOML  
+module TOML
   class Generator
     attr_reader :body, :doc
 
-    def initialize(doc)
+    def initialize(doc, preserve = false)
       # Ensure all the to_toml methods are injected into the base Ruby classes
       # used by TOML.
       self.class.inject!
-      
+
       @doc = doc
-      @body = doc.to_toml
-      
-      return @body
+      @body = doc.to_toml("", preserve)
+
+      @body
     end
-    
+
     # Whether or not the injections have already been done.
     @@injected = false
     # Inject to_toml methods into the Ruby classes used by TOML (booleans,
@@ -22,8 +21,8 @@ def initialize(doc)
     # if something doesn't have a to_toml method).
     def self.inject!
       return if @@injected
-      require 'toml/monkey_patch'
+      require "toml/monkey_patch"
       @@injected = true
     end
-  end#Generator
-end#TOML
+  end # Generator
+end # TOML
diff --git a/lib/toml/monkey_patch.rb b/lib/toml/monkey_patch.rb
@@ -1,26 +1,41 @@
 # Adds to_toml methods to base Ruby classes used by the generator.
 class Object
   def toml_table?
-    self.kind_of?(Hash)
+    is_a?(Hash)
   end
+
   def toml_table_array?
-    self.kind_of?(Array) && self.first.toml_table?
+    is_a?(Array) && first.toml_table?
   end
 end
+
 class Hash
-  def to_toml(path = "")
-    return "" if self.empty?
+  def to_toml(path = "", preserve = false)
+    return "" if empty?
 
     tables = {}
     values = {}
-    self.keys.sort.each do |key|
-      val = self[key]
-      if val.kind_of?(NilClass)
-        next
-      elsif val.toml_table? || val.toml_table_array?
-        tables[key] = val
-      else
-        values[key] = val
+    if preserve
+      keys.each do |key|
+        val = self[key]
+        if val.is_a?(NilClass)
+          next
+        elsif val.toml_table? || val.toml_table_array?
+          tables[key] = val
+        else
+          values[key] = val
+        end
+      end
+    else
+      keys.sort.each do |key|
+        val = self[key]
+        if val.is_a?(NilClass)
+          next
+        elsif val.toml_table? || val.toml_table_array?
+          tables[key] = val
+        else
+          values[key] = val
+        end
       end
     end
 
@@ -50,38 +65,44 @@ def to_toml(path = "")
     toml
   end
 end
+
 class Array
   def to_toml(path = "")
-    unless self.map(&:class).uniq.length == 1
+    unless map(&:class).uniq.length == 1
       raise "All array values must be the same type"
     end
 
-    if self.first.toml_table?
+    if first.toml_table?
       toml = ""
-      self.each do |val|
+      each do |val|
         toml << "\n[[#{path}]]\n"
         toml << val.to_toml(path)
       end
-      return toml
+      toml
     else
-      "[" + self.map {|v| v.to_toml(path) }.join(",") + "]"
+      "[" + map { |v| v.to_toml(path) }.join(",") + "]"
     end
   end
 end
+
 class TrueClass
-  def to_toml(path = ""); "true"; end
+  def to_toml(path = "") = "true"
 end
+
 class FalseClass
-  def to_toml(path = ""); "false"; end
+  def to_toml(path = "") = "false"
 end
+
 class String
-  def to_toml(path = ""); self.inspect; end
+  def to_toml(path = "") = inspect
 end
+
 class Numeric
-  def to_toml(path = ""); self.to_s; end
+  def to_toml(path = "") = to_s
 end
+
 class DateTime
   def to_toml(path = "")
-    self.rfc3339
+    rfc3339
   end
 end
diff --git a/test/test_generator.rb b/test/test_generator.rb
@@ -1,9 +1,8 @@
+require "rubygems"
+require "bundler/setup"
 
-require 'rubygems'
-require 'bundler/setup'
-
-require 'toml'
-require 'minitest/autorun'
+require "toml"
+require "minitest/autorun"
 
 class TestGenerator < MiniTest::Test
   def setup
@@ -22,7 +21,7 @@ def setup
           "name" => "second",
           "sub_table_array" => [
             {
-              "sub_name" => "sub first",
+              "sub_name" => "sub first"
             },
             {
               "sub_name" => "sub second"
@@ -39,7 +38,6 @@ def setup
       "date" => DateTime.now,
       "nil" => nil
     }
-
   end
 
   def test_generator
@@ -61,5 +59,10 @@ def test_generator
     doc.each do |key, val|
       assert_equal val, doc_parsed[key]
     end
+
+    # Test preserving the structure
+    body = TOML::Generator.new(doc, true).body
+    refute body.split.first != "integer", "The first entry in the generated TOML string is not 'integer' (as per @doc)."
+    assert_equal body.split.first, "integer"
   end
 end