Improved handling of trailing document comments

acceptance_tests/basic.sh

@@ -12,6 +12,62 @@ testBasicEvalRoundTrip() {
   assertEquals 123 "$X"
 }
 
+testBasicTrailingContent() {
+  cat >test-trailing.yml <<EOL
+test:
+# this comment will be removed
+EOL
+
+  read -r -d '' expected << EOM
+test:
+# this comment will be removed
+EOM
+  X=$(./yq test-trailing.yml -P)
+  assertEquals "$expected" "$X"
+}
+
+testBasicTrailingContent() {
+  cat >test-trailing.yml <<EOL
+test:
+# this comment will be removed
+EOL
+
+  read -r -d '' expected << EOM
+test:
+# hi
+EOM
+  X=$(./yq '. footComment = "hi"' test-trailing.yml)
+  assertEquals "$expected" "$X"
+}
+
+testBasicTrailingContentEvalAll() {
+  cat >test-trailing.yml <<EOL
+test:
+# this comment will be removed
+EOL
+
+  read -r -d '' expected << EOM
+test:
+# this comment will be removed
+EOM
+  X=$(./yq ea test-trailing.yml -P)
+  assertEquals "$expected" "$X"
+}
+
+testBasicTrailingContentEvalAll() {
+  cat >test-trailing.yml <<EOL
+test:
+# this comment will be removed
+EOL
+
+  read -r -d '' expected << EOM
+test:
+# hi
+EOM
+  X=$(./yq ea '. footComment = "hi"' test-trailing.yml)
+  assertEquals "$expected" "$X"
+}
+
 testBasicPipeWithDot() {
   ./yq -n ".a = 123" > test.yml
   X=$(cat test.yml | ./yq '.')

pkg/yqlib/candidate_node.go

@@ -14,7 +14,8 @@ type CandidateNode struct {
 	Parent *CandidateNode // parent node
 	Key    *yaml.Node     // node key, if this is a value from a map (or index in an array)
 
-	LeadingContent string
+	LeadingContent  string
+	TrailingContent string
 
 	Path      []interface{} /// the path we took to get to this node
 	Document  uint          // the document index of this node
@@ -170,6 +171,9 @@ func (n *CandidateNode) UpdateAttributesFrom(other *CandidateNode, prefs assignP
 	if other.Node.FootComment != "" {
 		n.Node.FootComment = other.Node.FootComment
 	}
+	if other.TrailingContent != "" {
+		n.TrailingContent = other.TrailingContent
+	}
 	if other.Node.HeadComment != "" {
 		n.Node.HeadComment = other.Node.HeadComment
 	}

pkg/yqlib/doc/operators/comment-operators.md

@@ -74,7 +74,6 @@ yq '. foot_comment=.a' sample.yml
 will output
 ```yaml
 a: cat
-
 # cat
 ```
 
@@ -117,6 +116,7 @@ b:
 Given a sample.yml file of:
 ```yaml
 a: cat # meow
+# have a great day
 ```
 then
 ```bash

pkg/yqlib/operator_comments.go

@@ -61,7 +61,9 @@ func assignCommentsOperator(d *dataTreeNavigator, context Context, expressionNod
 			candidate.Node.HeadComment = comment
 			candidate.LeadingContent = "" // clobber the leading content, if there was any.
 		}
-		if preferences.FootComment {
+		if preferences.FootComment && candidate.Node.Kind == yaml.DocumentNode && comment != "" {
+			candidate.TrailingContent = "# " + comment
+		} else if preferences.FootComment && candidate.Node.Kind != yaml.DocumentNode {
 			candidate.Node.FootComment = comment
 		}
 
@@ -97,6 +99,8 @@ func getCommentsOperator(d *dataTreeNavigator, context Context, expressionNode *
 			comment = chompRegexp.ReplaceAllString(comment, "")
 		} else if preferences.HeadComment {
 			comment = candidate.Node.HeadComment
+		} else if preferences.FootComment && candidate.Node.Kind == yaml.DocumentNode && candidate.TrailingContent != "" {
+			comment = candidate.TrailingContent
 		} else if preferences.FootComment {
 			comment = candidate.Node.FootComment
 		}

pkg/yqlib/operator_comments_test.go

@@ -59,7 +59,7 @@ var commentOperatorScenarios = []expressionScenario{
 		document:    `a: cat`,
 		expression:  `. foot_comment=.a`,
 		expected: []string{
-			"D0, P[], (doc)::a: cat\n\n# cat\n",
+			"D0, P[], (doc)::a: cat\n# cat\n",
 		},
 	},
 	{

pkg/yqlib/operator_multiply_test.go

@@ -144,7 +144,7 @@ var multiplyOperatorScenarios = []expressionScenario{
 		document2:  docNoComments,
 		expression: `select(fi == 0) * select(fi == 1)`,
 		expected: []string{
-			"D0, P[], (!!map)::a: apple\nb: banana\n\n# footer\n",
+			"D0, P[], (!!map)::a: apple\nb: banana\n# footer\n",
 		},
 	},
 	{
@@ -157,12 +157,13 @@ var multiplyOperatorScenarios = []expressionScenario{
 		},
 	},
 	{
-		skipDoc:    true,
-		document:   docNoComments,
-		document2:  docWithFooter,
-		expression: `select(fi == 0) * select(fi == 1)`,
+		description: "doc 2 has footer",
+		skipDoc:     true,
+		document:    docNoComments,
+		document2:   docWithFooter,
+		expression:  `select(fi == 0) * select(fi == 1)`,
 		expected: []string{
-			"D0, P[], (!!map)::b: banana\na: apple\n\n# footer\n",
+			"D0, P[], (!!map)::b: banana\na: apple\n# footer\n",
 		},
 	},
 	{

pkg/yqlib/printer.go

@@ -134,6 +134,10 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 			return err
 		}
 
+		if err := p.encoder.PrintLeadingContent(writer, mappedDoc.TrailingContent); err != nil {
+			return err
+		}
+
 		p.previousDocIndex = mappedDoc.Document
 		if err := writer.Flush(); err != nil {
 			return err

pkg/yqlib/stream_evaluator.go

@@ -110,6 +110,11 @@ func (s *streamEvaluator) Evaluate(filename string, reader io.Reader, node *Expr
 			Node:      &dataBucket,
 			FileIndex: s.fileIndex,
 		}
+		//move document comments into candidate node
+		// otherwise unwrap drops them.
+		candidateNode.TrailingContent = dataBucket.FootComment
+		dataBucket.FootComment = ""
+
 		if currentIndex == 0 {
 			candidateNode.LeadingContent = leadingContent
 		}

pkg/yqlib/utils.go

@@ -96,6 +96,11 @@ func readDocuments(reader io.Reader, filename string, fileIndex int, decoder Dec
 			EvaluateTogether: true,
 		}
 
+		//move document comments into candidate node
+		// otherwise unwrap drops them.
+		candidateNode.TrailingContent = dataBucket.FootComment
+		dataBucket.FootComment = ""
+
 		inputList.PushBack(candidateNode)
 
 		currentIndex = currentIndex + 1