testing implcit document ends

4 files changed, 32 insertions, 3 deletions

diff --git a/lib/psych/nodes/document.rb b/lib/psych/nodes/document.rb
index be786c6..4f46644 100644
--- a/lib/psych/nodes/document.rb
+++ b/lib/psych/nodes/document.rb
@@ -10,11 +10,15 @@ module Psych
       # Was this document implicitly created?
       attr_accessor :implicit
 
+      # Is the end of the document implicit?
+      attr_accessor :implicit_end
+
       def initialize version = [], tag_directives = [], implicit = false
         super()
         @version        = version
         @tag_directives = tag_directives
         @implicit       = implicit
+        @implicit_end   = true
       end
 
       ###
diff --git a/lib/psych/tree_builder.rb b/lib/psych/tree_builder.rb
index efbeaf9..089ee96 100644
--- a/lib/psych/tree_builder.rb
+++ b/lib/psych/tree_builder.rb
@@ -17,7 +17,6 @@ module Psych
     end
 
     %w{
-      Document
       Sequence
       Mapping
     }.each do |node|
@@ -28,12 +27,22 @@ module Psych
           @stack.push n
         end
 
-        def end_#{node.downcase}(*args)
+        def end_#{node.downcase}
           @stack.pop
         end
       }
     end
 
+    def start_document(*args)
+      n = Nodes::Document.new(*args)
+      @stack.last.children << n
+      @stack.push n
+    end
+
+    def end_document implicit_end
+      @stack.pop.implicit_end = implicit_end
+    end
+
     def start_stream encoding
       @stack.push Nodes::Stream.new(encoding)
     end
diff --git a/lib/psych/visitors/emitter.rb b/lib/psych/visitors/emitter.rb
index 71c33e3..966a09e 100644
--- a/lib/psych/visitors/emitter.rb
+++ b/lib/psych/visitors/emitter.rb
@@ -14,7 +14,7 @@ module Psych
       visitor_for(Nodes::Document) do |o|
         @handler.start_document o.version, o.tag_directives, o.implicit
         o.children.each { |c| c.accept self }
-        @handler.end_document o.implicit
+        @handler.end_document o.implicit_end
       end
 
       visitor_for(Nodes::Scalar) do |o|
diff --git a/test/visitors/test_emitter.rb b/test/visitors/test_emitter.rb
index cdca130..6f19df9 100644
--- a/test/visitors/test_emitter.rb
+++ b/test/visitors/test_emitter.rb
@@ -29,6 +29,22 @@ module Psych
         assert_equal @io.string, s.to_yaml
       end
 
+      def test_document_implicit_end
+        s       = Nodes::Stream.new
+        doc     = Nodes::Document.new
+        mapping = Nodes::Mapping.new
+        mapping.children << Nodes::Scalar.new('key')
+        mapping.children << Nodes::Scalar.new('value')
+        doc.children << mapping
+        s.children << doc
+
+        @visitor.accept s
+
+        assert_match(/key: value/, @io.string)
+        assert_equal @io.string, s.to_yaml
+        assert(/\.\.\./ !~ s.to_yaml)
+      end
+
       def test_scalar
         s       = Nodes::Stream.new
         doc     = Nodes::Document.new