Length (and some other operators) should not include head/foot comments #1231

examples/data1.yaml

@@ -1,2 +1,12 @@
-a: mike
-b: [t, f]
+# --------------------------------------------------
+# It's a test with comment
+# --------------------------------------------------
+test:
+  - name: a
+  - name: b
+  - name: c
+groups:
+  - name: a
+  - name: b
+  - name: c
+  - name: c

examples/data2.yaml

@@ -1,2 +1,5 @@
-c:
-  d: hamster
+# --------------------------------------------------
+# It's a test with comment
+# --------------------------------------------------
+groups:
+  - name: d

pkg/yqlib/candidate_node.go

@@ -86,18 +86,24 @@ func (n *CandidateNode) CreateChildInArray(index int, node *yaml.Node) *Candidat
 
 func (n *CandidateNode) CreateReplacement(node *yaml.Node) *CandidateNode {
 	return &CandidateNode{
-		Node:           node,
-		Path:           n.createChildPath(nil),
-		LeadingContent: n.LeadingContent,
-		Parent:         n.Parent,
-		Key:            n.Key,
-		IsMapKey:       n.IsMapKey,
-		Document:       n.Document,
-		Filename:       n.Filename,
-		FileIndex:      n.FileIndex,
+		Node:      node,
+		Path:      n.createChildPath(nil),
+		Parent:    n.Parent,
+		Key:       n.Key,
+		IsMapKey:  n.IsMapKey,
+		Document:  n.Document,
+		Filename:  n.Filename,
+		FileIndex: n.FileIndex,
 	}
 }
 
+func (n *CandidateNode) CreateReplacementWithDocWrappers(node *yaml.Node) *CandidateNode {
+	replacement := n.CreateReplacement(node)
+	replacement.LeadingContent = n.LeadingContent
+	replacement.TrailingContent = n.TrailingContent
+	return replacement
+}
+
 func (n *CandidateNode) createChildPath(path interface{}) []interface{} {
 	if path == nil {
 		newPath := make([]interface{}, len(n.Path))

pkg/yqlib/operator_collect_test.go

@@ -29,14 +29,23 @@ var collectOperatorScenarios = []expressionScenario{
 		},
 	},
 	{
-		skipDoc:  true,
-		document: "{a: apple}\n---\n{b: frog}",
-
+		skipDoc:    true,
+		document:   "{a: apple}\n---\n{b: frog}",
 		expression: `[.]`,
 		expected: []string{
 			"D0, P[], (!!seq)::- {a: apple}\n- {b: frog}\n",
 		},
 	},
+	{
+		description: "with comments",
+		skipDoc:     true,
+		document:    "# abc\n[{a: apple}]\n\n# xyz\n",
+
+		expression: `[.[]]`,
+		expected: []string{
+			"D0, P[], (!!seq)::- {a: apple}\n",
+		},
+	},
 	{
 		skipDoc:    true,
 		document:   ``,

pkg/yqlib/operator_entries.go

@@ -20,7 +20,7 @@ func entrySeqFor(key *yaml.Node, value *yaml.Node) *yaml.Node {
 
 func toEntriesFromMap(candidateNode *CandidateNode) *CandidateNode {
 	var sequence = &yaml.Node{Kind: yaml.SequenceNode, Tag: "!!seq"}
-	var entriesNode = candidateNode.CreateReplacement(sequence)
+	var entriesNode = candidateNode.CreateReplacementWithDocWrappers(sequence)
 
 	var contents = unwrapDoc(candidateNode.Node).Content
 	for index := 0; index < len(contents); index = index + 2 {
@@ -34,7 +34,7 @@ func toEntriesFromMap(candidateNode *CandidateNode) *CandidateNode {
 
 func toEntriesfromSeq(candidateNode *CandidateNode) *CandidateNode {
 	var sequence = &yaml.Node{Kind: yaml.SequenceNode, Tag: "!!seq"}
-	var entriesNode = candidateNode.CreateReplacement(sequence)
+	var entriesNode = candidateNode.CreateReplacementWithDocWrappers(sequence)
 
 	var contents = unwrapDoc(candidateNode.Node).Content
 	for index := 0; index < len(contents); index = index + 1 {
@@ -94,7 +94,7 @@ func parseEntry(entry *yaml.Node, position int) (*yaml.Node, *yaml.Node, error)
 
 func fromEntries(candidateNode *CandidateNode) (*CandidateNode, error) {
 	var node = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
-	var mapCandidateNode = candidateNode.CreateReplacement(node)
+	var mapCandidateNode = candidateNode.CreateReplacementWithDocWrappers(node)
 
 	var contents = unwrapDoc(candidateNode.Node).Content
 
@@ -141,9 +141,11 @@ func withEntriesOperator(d *dataTreeNavigator, context Context, expressionNode *
 	var results = list.New()
 
 	for el := toEntries.MatchingNodes.Front(); el != nil; el = el.Next() {
+		candidate := el.Value.(*CandidateNode)
+
 		//run expression against entries
 		// splat toEntries and pipe it into Rhs
-		splatted, err := splat(context.SingleChildContext(el.Value.(*CandidateNode)), traversePreferences{})
+		splatted, err := splat(context.SingleChildContext(candidate), traversePreferences{})
 		if err != nil {
 			return Context{}, err
 		}
@@ -160,6 +162,10 @@ func withEntriesOperator(d *dataTreeNavigator, context Context, expressionNode *
 		if err != nil {
 			return Context{}, err
 		}
+		collected.LeadingContent = candidate.LeadingContent
+		collected.TrailingContent = candidate.TrailingContent
+
+		log.Debugf("**** collected %v", collected.LeadingContent)
 
 		fromEntries, err := fromEntriesOperator(d, context.SingleChildContext(collected), expressionNode)
 		if err != nil {

pkg/yqlib/operator_entries_test.go

@@ -62,6 +62,15 @@ var entriesOperatorScenarios = []expressionScenario{
 			"D0, P[], (!!map)::KEY_c: 1\nKEY_d: 2\n",
 		},
 	},
+	{
+		skipDoc:    true,
+		document:   `[{a: 1, b: 2}, {c: 1, d: 2}]`,
+		expression: `.[] | with_entries(.key |= "KEY_" + .)`,
+		expected: []string{
+			"D0, P[], (!!map)::KEY_a: 1\nKEY_b: 2\n",
+			"D0, P[], (!!map)::KEY_c: 1\nKEY_d: 2\n",
+		},
+	},
 	{
 		description: "Use with_entries to filter the map",
 		document:    `{a: { b: bird }, c: { d: dog }}`,
@@ -70,6 +79,15 @@ var entriesOperatorScenarios = []expressionScenario{
 			"D0, P[], (!!map)::a: {b: bird}\n",
 		},
 	},
+	{
+		description: "Use with_entries to filter the map; head comment",
+		skipDoc:     true,
+		document:    "# abc\n{a: { b: bird }, c: { d: dog }}\n# xyz",
+		expression:  `with_entries(select(.value | has("b")))`,
+		expected: []string{
+			"D0, P[], (!!map)::# abc\na: {b: bird}\n# xyz\n",
+		},
+	},
 }
 
 func TestEntriesOperatorScenarios(t *testing.T) {

pkg/yqlib/operator_env_test.go

@@ -150,6 +150,16 @@ var envOperatorScenarios = []expressionScenario{
 		expression:    `"the ${notThere} ${alsoNotThere}" | envsubst(nu,ff)`,
 		expectedError: "variable ${notThere} not set",
 	},
+	{
+		description:          "with header/footer",
+		skipDoc:              true,
+		environmentVariables: map[string]string{"myenv": "cat meow"},
+		document:             "# abc\n{v: \"${myenv}\"}\n# xyz\n",
+		expression:           `(.. | select(tag == "!!str")) |= envsubst`,
+		expected: []string{
+			"D0, P[], (!!map)::# abc\n{v: \"cat meow\"}\n# xyz\n",
+		},
+	},
 }
 
 func TestEnvOperatorScenarios(t *testing.T) {

pkg/yqlib/operator_length_test.go

@@ -32,7 +32,7 @@ var lengthOperatorScenarios = []expressionScenario{
 	},
 	{
 		skipDoc:    true,
-		document:   `{a: key no exist}`,
+		document:   "# abc\n{a: key no exist}",
 		expression: `.b | length`,
 		expected: []string{
 			"D0, P[b], (!!int)::0\n",

pkg/yqlib/operator_pick.go

@@ -79,7 +79,7 @@ func pickOperator(d *dataTreeNavigator, context Context, expressionNode *Express
 			return Context{}, fmt.Errorf("cannot pick indicies from type %v (%v)", node.Tag, candidate.GetNicePath())
 		}
 
-		results.PushBack(candidate.CreateReplacement(replacement))
+		results.PushBack(candidate.CreateReplacementWithDocWrappers(replacement))
 	}
 
 	return context.ChildContext(results), nil

pkg/yqlib/operator_pick_test.go

@@ -14,6 +14,15 @@ var pickOperatorScenarios = []expressionScenario{
 			"D0, P[], (doc)::myMap: {hamster: squeek, cat: meow}\n",
 		},
 	},
+	{
+		description: "Pick keys from map with comments",
+		skipDoc:     true,
+		document:    "# abc\nmyMap: {cat: meow, dog: bark, thing: hamster, hamster: squeek}\n# xyz\n",
+		expression:  `.myMap |= pick(["hamster", "cat", "goat"])`,
+		expected: []string{
+			"D0, P[], (doc)::# abc\nmyMap: {hamster: squeek, cat: meow}\n# xyz\n",
+		},
+	},
 	{
 		description:    "Pick indices from array",
 		subdescription: "Note that the order of the indexes matches the pick order and non existent indexes are skipped.",
@@ -23,6 +32,15 @@ var pickOperatorScenarios = []expressionScenario{
 			"D0, P[], (!!seq)::[lion, cat]\n",
 		},
 	},
+	{
+		description: "Pick indices from array with comments",
+		skipDoc:     true,
+		document:    "# abc\n[cat, leopard, lion]\n# xyz",
+		expression:  `pick([2, 0, 734, -5])`,
+		expected: []string{
+			"D0, P[], (!!seq)::# abc\n[lion, cat]\n# xyz\n",
+		},
+	},
 }
 
 func TestPickOperatorScenarios(t *testing.T) {

pkg/yqlib/operator_reverse.go

@@ -25,7 +25,7 @@ func reverseOperator(d *dataTreeNavigator, context Context, expressionNode *Expr
 		for i, originalNode := range candidateNode.Content {
 			reverseList.Content[len(candidateNode.Content)-i-1] = originalNode
 		}
-		results.PushBack(candidate.CreateReplacement(reverseList))
+		results.PushBack(candidate.CreateReplacementWithDocWrappers(reverseList))
 
 	}
 

pkg/yqlib/operator_reverse_test.go

@@ -44,6 +44,15 @@ var reverseOperatorScenarios = []expressionScenario{
 			"D0, P[], (!!seq)::[{a: cat}, {a: banana}, {a: apple}]\n",
 		},
 	},
+	{
+		description: "Sort descending by string field, with comments",
+		skipDoc:     true,
+		document:    "# abc\n[{a: banana},{a: cat},{a: apple}]\n# xyz",
+		expression:  `sort_by(.a) | reverse`,
+		expected: []string{
+			"D0, P[], (!!seq)::# abc\n[{a: cat}, {a: banana}, {a: apple}]\n# xyz\n",
+		},
+	},
 }
 
 func TestReverseOperatorScenarios(t *testing.T) {

pkg/yqlib/operator_sort.go

@@ -64,7 +64,7 @@ func sortByOperator(d *dataTreeNavigator, context Context, expressionNode *Expre
 		for i, sortedNode := range sortableArray {
 			sortedList.Content[i] = sortedNode.Node
 		}
-		results.PushBack(candidate.CreateReplacement(sortedList))
+		results.PushBack(candidate.CreateReplacementWithDocWrappers(sortedList))
 	}
 	return context.ChildContext(results), nil
 }

pkg/yqlib/operator_unique.go

@@ -56,7 +56,7 @@ func uniqueBy(d *dataTreeNavigator, context Context, expressionNode *ExpressionN
 			resultNode.Content = append(resultNode.Content, el.Value.(*yaml.Node))
 		}
 
-		results.PushBack(candidate.CreateReplacement(resultNode))
+		results.PushBack(candidate.CreateReplacementWithDocWrappers(resultNode))
 	}
 
 	return context.ChildContext(results), nil

pkg/yqlib/operator_unique_test.go

@@ -56,6 +56,14 @@ var uniqueOperatorScenarios = []expressionScenario{
 			"D0, P[], (!!seq)::- {name: harry, pet: cat}\n",
 		},
 	},
+	{
+		skipDoc:    true,
+		document:   "# abc\n[{name: harry, pet: cat}, {pet: fish}, {name: harry, pet: dog}]\n# xyz",
+		expression: `unique_by(.name)`,
+		expected: []string{
+			"D0, P[], (!!seq)::# abc\n- {name: harry, pet: cat}\n- {pet: fish}\n# xyz\n",
+		},
+	},
 }
 
 func TestUniqueOperatorScenarios(t *testing.T) {