Preserve comments on map keys

pkg/yqlib/candidate_node.go

@@ -2,8 +2,6 @@ package yqlib
 
 import (
 	"fmt"
-	"strconv"
-	"strings"
 
 	"github.com/jinzhu/copier"
 	yaml "gopkg.in/yaml.v3"
@@ -18,10 +16,15 @@ type CandidateNode struct {
 	// when performing op against all nodes given, this will treat all the nodes as one
 	// (e.g. top level cross document merge). This property does not propegate to child nodes.
 	EvaluateTogether bool
+	IsMapKey         bool
 }
 
 func (n *CandidateNode) GetKey() string {
-	return fmt.Sprintf("%v - %v", n.Document, n.Path)
+	keyPrefix := ""
+	if n.IsMapKey {
+		keyPrefix = "key-"
+	}
+	return fmt.Sprintf("%v%v - %v", keyPrefix, n.Document, n.Path)
 }
 
 func (n *CandidateNode) CreateChild(path interface{}, node *yaml.Node) *CandidateNode {
@@ -66,6 +69,7 @@ func (n *CandidateNode) UpdateFrom(other *CandidateNode) {
 }
 
 func (n *CandidateNode) UpdateAttributesFrom(other *CandidateNode) {
+	log.Debug("UpdateAttributesFrom: n: %v other: %v", n.GetKey(), other.GetKey())
 	if n.Node.Kind != other.Node.Kind {
 		// clear out the contents when switching to a different type
 		// e.g. map to array
@@ -86,46 +90,3 @@ func (n *CandidateNode) UpdateAttributesFrom(other *CandidateNode) {
 	n.Node.HeadComment = n.Node.HeadComment + other.Node.HeadComment
 	n.Node.LineComment = n.Node.LineComment + other.Node.LineComment
 }
-
-func (n *CandidateNode) PathStackToString() string {
-	return mergePathStackToString(n.Path)
-}
-
-func mergePathStackToString(pathStack []interface{}) string {
-	var sb strings.Builder
-	for index, path := range pathStack {
-		switch path.(type) {
-		case int, int64:
-			// if arrayMergeStrategy == AppendArrayMergeStrategy {
-			// sb.WriteString("[+]")
-			// } else {
-			sb.WriteString(fmt.Sprintf("[%v]", path))
-			// }
-
-		default:
-			s := fmt.Sprintf("%v", path)
-			var _, errParsingInt = strconv.ParseInt(s, 10, 64) // nolint
-
-			hasSpecial := strings.Contains(s, ".") || strings.Contains(s, "[") || strings.Contains(s, "]") || strings.Contains(s, "\"")
-			hasDoubleQuotes := strings.Contains(s, "\"")
-			wrappingCharacterStart := "\""
-			wrappingCharacterEnd := "\""
-			if hasDoubleQuotes {
-				wrappingCharacterStart = "("
-				wrappingCharacterEnd = ")"
-			}
-			if hasSpecial || errParsingInt == nil {
-				sb.WriteString(wrappingCharacterStart)
-			}
-			sb.WriteString(s)
-			if hasSpecial || errParsingInt == nil {
-				sb.WriteString(wrappingCharacterEnd)
-			}
-		}
-
-		if index < len(pathStack)-1 {
-			sb.WriteString(".")
-		}
-	}
-	return sb.String()
-}

pkg/yqlib/doc/Multiply.md

@@ -71,7 +71,6 @@ Given a sample.yml file of:
 a: {things: great}
 b:
   also: "me"
-
 ```
 then
 ```bash

pkg/yqlib/operator_assign.go

@@ -36,6 +36,7 @@ func assignUpdateOperator(d *dataTreeNavigator, context Context, expressionNode
 
 // does not update content or values
 func assignAttributesOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
+	log.Debug("getting lhs matching nodes for update")
 	lhs, err := d.GetMatchingNodes(context, expressionNode.Lhs)
 	if err != nil {
 		return Context{}, err

pkg/yqlib/operator_delete.go

@@ -26,7 +26,7 @@ func deleteChildOperator(d *dataTreeNavigator, context Context, expressionNode *
 
 			deleteImmediateChildOpNode := &ExpressionNode{
 				Operation: deleteImmediateChildOp,
-				Rhs:       createTraversalTree(candidate.Path[0:len(candidate.Path)-1], traversePreferences{}),
+				Rhs:       createTraversalTree(candidate.Path[0:len(candidate.Path)-1], traversePreferences{}, false),
 			}
 
 			_, err := d.GetMatchingNodes(contextToUse, deleteImmediateChildOpNode)

pkg/yqlib/operator_multiply.go

@@ -66,7 +66,7 @@ func mergeObjects(d *dataTreeNavigator, context Context, lhs *CandidateNode, rhs
 
 	// shouldn't recurse arrays if appending
 	prefs := recursiveDescentPreferences{RecurseArray: !shouldAppendArrays,
-		TraversePreferences: traversePreferences{DontFollowAlias: true}}
+		TraversePreferences: traversePreferences{DontFollowAlias: true, IncludeMapKeys: true}}
 	err := recursiveDecent(d, results, context.SingleChildContext(rhs), prefs)
 	if err != nil {
 		return nil, err
@@ -78,7 +78,12 @@ func mergeObjects(d *dataTreeNavigator, context Context, lhs *CandidateNode, rhs
 	}
 
 	for el := results.Front(); el != nil; el = el.Next() {
-		err := applyAssignment(d, context, pathIndexToStartFrom, lhs, el.Value.(*CandidateNode), preferences)
+		candidate := el.Value.(*CandidateNode)
+		if candidate.Node.Tag == "!!merge" {
+			continue
+		}
+
+		err := applyAssignment(d, context, pathIndexToStartFrom, lhs, candidate, preferences)
 		if err != nil {
 			return nil, err
 		}
@@ -88,7 +93,7 @@ func mergeObjects(d *dataTreeNavigator, context Context, lhs *CandidateNode, rhs
 
 func applyAssignment(d *dataTreeNavigator, context Context, pathIndexToStartFrom int, lhs *CandidateNode, rhs *CandidateNode, preferences multiplyPreferences) error {
 	shouldAppendArrays := preferences.AppendArrays
-	log.Debugf("merge - applyAssignment lhs %v, rhs: %v", NodeToString(lhs), NodeToString(rhs))
+	log.Debugf("merge - applyAssignment lhs %v, rhs: %v", lhs.GetKey(), rhs.GetKey())
 
 	lhsPath := rhs.Path[pathIndexToStartFrom:]
 
@@ -101,7 +106,7 @@ func applyAssignment(d *dataTreeNavigator, context Context, pathIndexToStartFrom
 	}
 	rhsOp := &Operation{OperationType: valueOpType, CandidateNode: rhs}
 
-	assignmentOpNode := &ExpressionNode{Operation: assignmentOp, Lhs: createTraversalTree(lhsPath, preferences.TraversePrefs), Rhs: &ExpressionNode{Operation: rhsOp}}
+	assignmentOpNode := &ExpressionNode{Operation: assignmentOp, Lhs: createTraversalTree(lhsPath, preferences.TraversePrefs, rhs.IsMapKey), Rhs: &ExpressionNode{Operation: rhsOp}}
 
 	_, err := d.GetMatchingNodes(context.SingleChildContext(lhs), assignmentOpNode)
 

pkg/yqlib/operator_multiply_test.go

@@ -13,6 +13,22 @@ var multiplyOperatorScenarios = []expressionScenario{
 			"D0, P[], (!!map)::{a: {also: me}, b: {also: me}}\n",
 		},
 	},
+	{
+		skipDoc:    true,
+		document:   "# b\nb:\n  # a\n  a: cat",
+		expression: "{} * .",
+		expected: []string{
+			"D0, P[], (!!map)::# b\nb:\n    # a\n    a: cat\n",
+		},
+	},
+	{
+		skipDoc:    true,
+		document:   "# b\nb:\n  # a\n  a: cat",
+		expression: ". * {}",
+		expected: []string{
+			"D0, P[], (!!map)::# b\nb:\n    # a\n    a: cat\n",
+		},
+	},
 	{
 		skipDoc:    true,
 		document:   `{a: &a { b: &b { c: &c cat } } }`,
@@ -100,7 +116,7 @@ var multiplyOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{a: {things: great}, b: {also: me}}`,
-		expression: `. * {"a":.b}`,
+		expression: `. * {"a": .b}`,
 		expected: []string{
 			"D0, P[], (!!map)::{a: {things: great, also: me}, b: {also: me}}\n",
 		},
@@ -108,16 +124,10 @@ var multiplyOperatorScenarios = []expressionScenario{
 	{
 		description:           "Merge keeps style of LHS",
 		dontFormatInputForDoc: true,
-		document: `a: {things: great}
-b:
-  also: "me"
-`,
-		expression: `. * {"a":.b}`,
+		document:              "a: {things: great}\nb:\n  also: \"me\"",
+		expression:            `. * {"a":.b}`,
 		expected: []string{
-			`D0, P[], (!!map)::a: {things: great, also: "me"}
-b:
-    also: "me"
-`,
+			"D0, P[], (!!map)::a: {things: great, also: \"me\"}\nb:\n    also: \"me\"\n",
 		},
 	},
 	{

pkg/yqlib/operator_traverse_path.go

@@ -10,9 +10,10 @@ import (
 )
 
 type traversePreferences struct {
-	DontFollowAlias bool
-	IncludeMapKeys  bool
-	DontAutoCreate  bool // by default, we automatically create entries on the fly.
+	DontFollowAlias      bool
+	IncludeMapKeys       bool
+	DontAutoCreate       bool // by default, we automatically create entries on the fly.
+	DontIncludeMapValues bool
 }
 
 func splat(d *dataTreeNavigator, context Context, prefs traversePreferences) (Context, error) {
@@ -214,10 +215,21 @@ func traverseMap(context Context, matchingNode *CandidateNode, key string, prefs
 	if !prefs.DontAutoCreate && !context.DontAutoCreate && newMatches.Len() == 0 {
 		//no matches, create one automagically
 		valueNode := &yaml.Node{Tag: "!!null", Kind: yaml.ScalarNode, Value: "null"}
+		keyNode := &yaml.Node{Kind: yaml.ScalarNode, Value: key}
 		node := matchingNode.Node
-		node.Content = append(node.Content, &yaml.Node{Kind: yaml.ScalarNode, Value: key}, valueNode)
-		candidateNode := matchingNode.CreateChild(key, valueNode)
-		newMatches.Set(candidateNode.GetKey(), candidateNode)
+		node.Content = append(node.Content, keyNode, valueNode)
+
+		if prefs.IncludeMapKeys {
+			log.Debug("including key")
+			candidateNode := matchingNode.CreateChild(key, keyNode)
+			candidateNode.IsMapKey = true
+			newMatches.Set(fmt.Sprintf("keyOf-%v", candidateNode.GetKey()), candidateNode)
+		}
+		if !prefs.DontIncludeMapValues {
+			log.Debug("including value")
+			candidateNode := matchingNode.CreateChild(key, valueNode)
+			newMatches.Set(candidateNode.GetKey(), candidateNode)
+		}
 	}
 
 	results := list.New()
@@ -253,11 +265,16 @@ func doTraverseMap(newMatches *orderedmap.OrderedMap, candidate *CandidateNode,
 		} else if splat || keyMatches(key, wantedKey) {
 			log.Debug("MATCHED")
 			if prefs.IncludeMapKeys {
+				log.Debug("including key")
 				candidateNode := candidate.CreateChild(key.Value, key)
+				candidateNode.IsMapKey = true
 				newMatches.Set(fmt.Sprintf("keyOf-%v", candidateNode.GetKey()), candidateNode)
 			}
-			candidateNode := candidate.CreateChild(key.Value, value)
-			newMatches.Set(candidateNode.GetKey(), candidateNode)
+			if !prefs.DontIncludeMapValues {
+				log.Debug("including value")
+				candidateNode := candidate.CreateChild(key.Value, value)
+				newMatches.Set(candidateNode.GetKey(), candidateNode)
+			}
 		}
 	}
 

pkg/yqlib/operators.go

@@ -4,6 +4,7 @@ import (
 	"container/list"
 	"fmt"
 
+	"github.com/jinzhu/copier"
 	"gopkg.in/yaml.v3"
 )
 
@@ -88,16 +89,25 @@ func createBooleanCandidate(owner *CandidateNode, value bool) *CandidateNode {
 	return owner.CreateChild(nil, node)
 }
 
-func createTraversalTree(path []interface{}, traversePrefs traversePreferences) *ExpressionNode {
+func createTraversalTree(path []interface{}, traversePrefs traversePreferences, targetKey bool) *ExpressionNode {
 	if len(path) == 0 {
 		return &ExpressionNode{Operation: &Operation{OperationType: selfReferenceOpType}}
 	} else if len(path) == 1 {
-		return &ExpressionNode{Operation: &Operation{OperationType: traversePathOpType, Preferences: traversePrefs, Value: path[0], StringValue: fmt.Sprintf("%v", path[0])}}
+		lastPrefs := traversePrefs
+		if targetKey {
+			err := copier.Copy(&lastPrefs, traversePrefs)
+			if err != nil {
+				panic(err)
+			}
+			lastPrefs.IncludeMapKeys = true
+			lastPrefs.DontIncludeMapValues = true
+		}
+		return &ExpressionNode{Operation: &Operation{OperationType: traversePathOpType, Preferences: lastPrefs, Value: path[0], StringValue: fmt.Sprintf("%v", path[0])}}
 	}
 
 	return &ExpressionNode{
 		Operation: &Operation{OperationType: shortPipeOpType},
-		Lhs:       createTraversalTree(path[0:1], traversePrefs),
-		Rhs:       createTraversalTree(path[1:], traversePrefs),
+		Lhs:       createTraversalTree(path[0:1], traversePrefs, false),
+		Rhs:       createTraversalTree(path[1:], traversePrefs, targetKey),
 	}
 }