Can assign-update comments

pkg/yqlib/doc/Comment Operators.md

@@ -13,6 +13,22 @@ will output
 a: cat # single
 ```
 
+## Use update assign to perform relative updates
+Given a sample.yml file of:
+```yaml
+a: cat
+b: dog
+```
+then
+```bash
+yq eval '.. lineComment |= .' sample.yml
+```
+will output
+```yaml
+a: cat # cat
+b: dog # dog
+```
+
 ## Set head comment
 Given a sample.yml file of:
 ```yaml

pkg/yqlib/lib.go

@@ -86,6 +86,7 @@ type Operation struct {
 	StringValue   string
 	CandidateNode *CandidateNode // used for Value Path elements
 	Preferences   interface{}
+	UpdateAssign  bool // used for assign ops, when true it means we evaluate the rhs given the lhs (instead of matching nodes)
 }
 
 func CreateValueOperation(value interface{}, stringValue string) *Operation {

pkg/yqlib/operator_add.go

@@ -16,7 +16,7 @@ func createSelfAddOp(rhs *PathTreeNode) *PathTreeNode {
 
 func AddAssignOperator(d *dataTreeNavigator, matchingNodes *list.List, pathNode *PathTreeNode) (*list.List, error) {
 	assignmentOp := &Operation{OperationType: Assign}
-	assignmentOp.Preferences = &AssignOpPreferences{true}
+	assignmentOp.UpdateAssign = true
 
 	assignmentOpNode := &PathTreeNode{Operation: assignmentOp, Lhs: pathNode.Lhs, Rhs: createSelfAddOp(pathNode.Rhs)}
 	return d.GetMatchingNodes(matchingNodes, assignmentOpNode)

pkg/yqlib/operator_assign.go

@@ -2,26 +2,20 @@ package yqlib
 
 import "container/list"
 
-type AssignOpPreferences struct {
-	UpdateAssign bool
-}
-
 func AssignUpdateOperator(d *dataTreeNavigator, matchingNodes *list.List, pathNode *PathTreeNode) (*list.List, error) {
 	lhs, err := d.GetMatchingNodes(matchingNodes, pathNode.Lhs)
 	if err != nil {
 		return nil, err
 	}
-	preferences := pathNode.Operation.Preferences.(*AssignOpPreferences)
-
 	var rhs *list.List
-	if !preferences.UpdateAssign {
+	if !pathNode.Operation.UpdateAssign {
 		rhs, err = d.GetMatchingNodes(matchingNodes, pathNode.Rhs)
 	}
 
 	for el := lhs.Front(); el != nil; el = el.Next() {
 		candidate := el.Value.(*CandidateNode)
 
-		if preferences.UpdateAssign {
+		if pathNode.Operation.UpdateAssign {
 			rhs, err = d.GetMatchingNodes(nodeToMap(candidate), pathNode.Rhs)
 		}
 

pkg/yqlib/operator_comments.go

@@ -4,7 +4,7 @@ import (
 	"container/list"
 	"strings"
 
-	"gopkg.in/yaml.v3"
+	yaml "gopkg.in/yaml.v3"
 )
 
 type CommentOpPreferences struct {
@@ -17,15 +17,6 @@ func AssignCommentsOperator(d *dataTreeNavigator, matchingNodes *list.List, path
 
 	log.Debugf("AssignComments operator!")
 
-	rhs, err := d.GetMatchingNodes(matchingNodes, pathNode.Rhs)
-	if err != nil {
-		return nil, err
-	}
-	comment := ""
-	if rhs.Front() != nil {
-		comment = rhs.Front().Value.(*CandidateNode).Node.Value
-	}
-
 	lhs, err := d.GetMatchingNodes(matchingNodes, pathNode.Lhs)
 
 	if err != nil {
@@ -34,8 +25,32 @@ func AssignCommentsOperator(d *dataTreeNavigator, matchingNodes *list.List, path
 
 	preferences := pathNode.Operation.Preferences.(*CommentOpPreferences)
 
+	comment := ""
+	if !pathNode.Operation.UpdateAssign {
+		rhs, err := d.GetMatchingNodes(matchingNodes, pathNode.Rhs)
+		if err != nil {
+			return nil, err
+		}
+
+		if rhs.Front() != nil {
+			comment = rhs.Front().Value.(*CandidateNode).Node.Value
+		}
+	}
+
 	for el := lhs.Front(); el != nil; el = el.Next() {
 		candidate := el.Value.(*CandidateNode)
+
+		if pathNode.Operation.UpdateAssign {
+			rhs, err := d.GetMatchingNodes(nodeToMap(candidate), pathNode.Rhs)
+			if err != nil {
+				return nil, err
+			}
+
+			if rhs.Front() != nil {
+				comment = rhs.Front().Value.(*CandidateNode).Node.Value
+			}
+		}
+
 		log.Debugf("Setting comment of : %v", candidate.GetKey())
 		if preferences.LineComment {
 			candidate.Node.LineComment = comment

pkg/yqlib/operator_comments_test.go

@@ -13,6 +13,39 @@ var commentOperatorScenarios = []expressionScenario{
 			"D0, P[], (doc)::a: cat # single\n",
 		},
 	},
+	{
+		skipDoc:    true,
+		document:   "a: cat\nb: dog",
+		expression: `.a lineComment=.b`,
+		expected: []string{
+			"D0, P[], (doc)::a: cat # dog\nb: dog\n",
+		},
+	},
+	{
+		skipDoc:    true,
+		document:   "a: cat\n---\na: dog",
+		expression: `.a lineComment |= documentIndex`,
+		expected: []string{
+			"D0, P[], (doc)::a: cat # 0\n",
+			"D1, P[], (doc)::a: dog # 1\n",
+		},
+	},
+	{
+		description: "Use update assign to perform relative updates",
+		document:    "a: cat\nb: dog",
+		expression:  `.. lineComment |= .`,
+		expected: []string{
+			"D0, P[], (!!map)::a: cat # cat\nb: dog # dog\n",
+		},
+	},
+	{
+		skipDoc:    true,
+		document:   "a: cat\nb: dog",
+		expression: `.. comments |= .`,
+		expected: []string{
+			"D0, P[], (!!map)::a: cat # cat\n# cat\n\n# cat\nb: dog # dog\n# dog\n\n# dog\n",
+		},
+	},
 	{
 		description: "Set head comment",
 		document:    `a: cat`,

pkg/yqlib/operator_multiply.go

@@ -115,7 +115,7 @@ func applyAssignment(d *dataTreeNavigator, pathIndexToStartFrom int, lhs *Candid
 	assignmentOp := &Operation{OperationType: AssignAttributes}
 	if rhs.Node.Kind == yaml.ScalarNode || rhs.Node.Kind == yaml.AliasNode {
 		assignmentOp.OperationType = Assign
-		assignmentOp.Preferences = &AssignOpPreferences{false}
+		assignmentOp.UpdateAssign = false
 	} else if shouldAppendArrays && rhs.Node.Kind == yaml.SequenceNode {
 		assignmentOp.OperationType = AddAssign
 	}

pkg/yqlib/path_parse_test.go

@@ -137,6 +137,11 @@ var pathTests = []struct {
 		append(make([]interface{}, 0), "SELF", "ASSIGN_COMMENT", "str (string)"),
 		append(make([]interface{}, 0), "SELF", "str (string)", "ASSIGN_COMMENT"),
 	},
+	{
+		`. lineComment |= "str"`,
+		append(make([]interface{}, 0), "SELF", "ASSIGN_COMMENT", "str (string)"),
+		append(make([]interface{}, 0), "SELF", "str (string)", "ASSIGN_COMMENT"),
+	},
 	{
 		`.a.b tag="!!str"`,
 		append(make([]interface{}, 0), "a", "SHORT_PIPE", "b", "ASSIGN_TAG", "!!str (string)"),

pkg/yqlib/path_tokeniser.go

@@ -92,6 +92,15 @@ func opAssignableToken(opType *OperationType, assignOpType *OperationType) lex.A
 	return opTokenWithPrefs(opType, assignOpType, nil)
 }
 
+func assignOpToken(updateAssign bool) lex.Action {
+	return func(s *lex.Scanner, m *machines.Match) (interface{}, error) {
+		log.Debug("assignOpToken %v", string(m.Bytes))
+		value := string(m.Bytes)
+		op := &Operation{OperationType: Assign, Value: Assign.Type, StringValue: value, UpdateAssign: updateAssign}
+		return &Token{TokenType: OperationToken, Operation: op}, nil
+	}
+}
+
 func opTokenWithPrefs(op *OperationType, assignOpType *OperationType, preferences interface{}) lex.Action {
 	return func(s *lex.Scanner, m *machines.Match) (interface{}, error) {
 		log.Debug("opTokenWithPrefs %v", string(m.Bytes))
@@ -105,6 +114,21 @@ func opTokenWithPrefs(op *OperationType, assignOpType *OperationType, preference
 	}
 }
 
+func assignAllCommentsOp(updateAssign bool) lex.Action {
+	return func(s *lex.Scanner, m *machines.Match) (interface{}, error) {
+		log.Debug("assignAllCommentsOp %v", string(m.Bytes))
+		value := string(m.Bytes)
+		op := &Operation{
+			OperationType: AssignComment,
+			Value:         AssignComment.Type,
+			StringValue:   value,
+			UpdateAssign:  updateAssign,
+			Preferences:   &CommentOpPreferences{LineComment: true, HeadComment: true, FootComment: true},
+		}
+		return &Token{TokenType: OperationToken, Operation: op}, nil
+	}
+}
+
 func literalToken(pType TokenType, checkForPost bool) lex.Action {
 	return func(s *lex.Scanner, m *machines.Match) (interface{}, error) {
 		return &Token{TokenType: pType, CheckForPostTraverse: checkForPost}, nil
@@ -210,16 +234,17 @@ func initLexer() (*lex.Lexer, error) {
 
 	lexer.Add([]byte(`footComment`), opTokenWithPrefs(GetComment, AssignComment, &CommentOpPreferences{FootComment: true}))
 
-	lexer.Add([]byte(`comments\s*=`), opTokenWithPrefs(AssignComment, nil, &CommentOpPreferences{LineComment: true, HeadComment: true, FootComment: true}))
+	lexer.Add([]byte(`comments\s*=`), assignAllCommentsOp(false))
+	lexer.Add([]byte(`comments\s*\|=`), assignAllCommentsOp(true))
 
 	lexer.Add([]byte(`collect`), opToken(Collect))
 
 	lexer.Add([]byte(`\s*==\s*`), opToken(Equals))
-	lexer.Add([]byte(`\s*=\s*`), opTokenWithPrefs(Assign, nil, &AssignOpPreferences{false}))
+	lexer.Add([]byte(`\s*=\s*`), assignOpToken(false))
 
 	lexer.Add([]byte(`del`), opToken(DeleteChild))
 
-	lexer.Add([]byte(`\s*\|=\s*`), opTokenWithPrefs(Assign, nil, &AssignOpPreferences{true}))
+	lexer.Add([]byte(`\s*\|=\s*`), assignOpToken(true))
 
 	lexer.Add([]byte("( |\t|\n|\r)+"), skip)
 
@@ -326,6 +351,7 @@ func (p *pathTokeniser) handleToken(tokens []*Token, index int, postProcessedTok
 		tokens[index+1].TokenType == OperationToken &&
 		tokens[index+1].Operation.OperationType == Assign {
 		token.Operation = token.AssignOperation
+		token.Operation.UpdateAssign = tokens[index+1].Operation.UpdateAssign
 		skipNextToken = true
 	}