assign operator

pkg/yqlib/treeops/data_tree_navigator.go

@@ -27,6 +27,7 @@ func NewDataTreeNavigator(navigationPrefs NavigationPrefs) DataTreeNavigator {
 	operatorHandlers[Equals] = EqualsOperator
 	operatorHandlers[Or] = UnionOperator
 	operatorHandlers[And] = IntersectionOperator
+	operatorHandlers[Assign] = AssignOperator
 
 	return &dataTreeNavigator{leafTraverser, operatorHandlers}
 }

pkg/yqlib/treeops/data_tree_navigator_test.go

@@ -78,6 +78,31 @@ func TestDataTreeNavigatorArraySimple(t *testing.T) {
 	test.AssertResult(t, expected, resultsToString(results))
 }
 
+func TestDataTreeNavigatorSimpleAssign(t *testing.T) {
+
+	nodes := readDoc(t, `a: 
+  b: apple`)
+
+	path, errPath := treeCreator.ParsePath("a.b := frog")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: [a b]
+  Tag: !!str, Kind: ScalarNode, Anchor: 
+  frog
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
 func TestDataTreeNavigatorArraySplat(t *testing.T) {
 
 	nodes := readDoc(t, `- b: apple

pkg/yqlib/treeops/operators.go

@@ -12,6 +12,18 @@ func TraverseOperator(d *dataTreeNavigator, matchingNodes *orderedmap.OrderedMap
 	return d.getMatchingNodes(lhs, pathNode.Rhs)
 }
 
+func AssignOperator(d *dataTreeNavigator, matchingNodes *orderedmap.OrderedMap, pathNode *PathTreeNode) (*orderedmap.OrderedMap, error) {
+	lhs, err := d.getMatchingNodes(matchingNodes, pathNode.Lhs)
+	if err != nil {
+		return nil, err
+	}
+	for el := lhs.Front(); el != nil; el = el.Next() {
+		node := el.Value.(*CandidateNode)
+		node.Node.Value = pathNode.Rhs.PathElement.StringValue
+	}
+	return lhs, nil
+}
+
 func UnionOperator(d *dataTreeNavigator, matchingNodes *orderedmap.OrderedMap, pathNode *PathTreeNode) (*orderedmap.OrderedMap, error) {
 	lhs, err := d.getMatchingNodes(matchingNodes, pathNode.Lhs)
 	if err != nil {

pkg/yqlib/treeops/path_postfix.go

@@ -26,6 +26,7 @@ const (
 	And
 	Equals
 	EqualsSelf
+	Assign
 )
 
 type PathElement struct {
@@ -54,6 +55,8 @@ func (p *PathElement) toString() string {
 			result = result + "EQUALS\n"
 		case EqualsSelf:
 			result = result + "EQUALS SELF\n"
+		case Assign:
+			result = result + "ASSIGN\n"
 		case Traverse:
 			result = result + "TRAVERSE\n"
 		}
@@ -82,6 +85,9 @@ func initMaps() {
 	precedenceMap[TokenIds["EQUALS_SELF_OPERATOR"]] = 30
 	operationTypeMapper[TokenIds["EQUALS_SELF_OPERATOR"]] = EqualsSelf
 
+	precedenceMap[TokenIds["ASSIGN_OPERATOR"]] = 35
+	operationTypeMapper[TokenIds["ASSIGN_OPERATOR"]] = Assign
+
 	precedenceMap[TokenIds["TRAVERSE_OPERATOR"]] = 40
 	operationTypeMapper[TokenIds["TRAVERSE_OPERATOR"]] = Traverse
 }
@@ -122,14 +128,6 @@ func (p *pathPostFixer) ConvertToPostfix(infixTokens []*lex.Token) ([]*PathEleme
 			result = append(result, &pathElement)
 		case TokenIds["("]:
 			opStack = append(opStack, token)
-		case TokenIds["OR_OPERATOR"], TokenIds["AND_OPERATOR"], TokenIds["EQUALS_OPERATOR"], TokenIds["EQUALS_SELF_OPERATOR"], TokenIds["TRAVERSE_OPERATOR"]:
-			var currentPrecedence = precedenceMap[token.Type]
-			// pop off higher precedent operators onto the result
-			for len(opStack) > 0 && precedenceMap[opStack[len(opStack)-1].Type] >= currentPrecedence {
-				opStack, result = popOpToResult(opStack, result)
-			}
-			// add this operator to the opStack
-			opStack = append(opStack, token)
 		case TokenIds[")"]:
 			for len(opStack) > 0 && opStack[len(opStack)-1].Type != TokenIds["("] {
 				opStack, result = popOpToResult(opStack, result)
@@ -139,6 +137,14 @@ func (p *pathPostFixer) ConvertToPostfix(infixTokens []*lex.Token) ([]*PathEleme
 			}
 			// now we should have ( as the last element on the opStack, get rid of it
 			opStack = opStack[0 : len(opStack)-1]
+		default:
+			var currentPrecedence = precedenceMap[token.Type]
+			// pop off higher precedent operators onto the result
+			for len(opStack) > 0 && precedenceMap[opStack[len(opStack)-1].Type] >= currentPrecedence {
+				opStack, result = popOpToResult(opStack, result)
+			}
+			// add this operator to the opStack
+			opStack = append(opStack, token)
 		}
 	}
 	return result, nil

pkg/yqlib/treeops/path_postfix_test.go

@@ -75,6 +75,28 @@ Operation - TRAVERSE
 	test.AssertResultComplex(t, expectedOutput, actual)
 }
 
+func TestPostFixSimpleAssign(t *testing.T) {
+	var infix = "a.b := frog"
+	var expectedOutput = `PathKey - 'a'
+--------
+PathKey - 'b'
+--------
+Operation - TRAVERSE
+--------
+PathKey - 'frog'
+--------
+Operation - ASSIGN
+--------
+`
+
+	actual, err := testExpression(infix)
+	if err != nil {
+		t.Error(err)
+	}
+
+	test.AssertResultComplex(t, expectedOutput, actual)
+}
+
 func TestPostFixSimplePathNumbersExample(t *testing.T) {
 	var infix = "apples[0].cat"
 	var expectedOutput = `PathKey - 'apples'

pkg/yqlib/treeops/path_tokeniser.go

@@ -27,6 +27,7 @@ func initTokens() {
 		"AND_OPERATOR",
 		"EQUALS_OPERATOR",
 		"EQUALS_SELF_OPERATOR",
+		"ASSIGN_OPERATOR",
 		"TRAVERSE_OPERATOR",
 		"PATH_KEY",    // apples
 		"ARRAY_INDEX", // 123
@@ -90,6 +91,7 @@ func initLexer() (*lex.Lexer, error) {
 	lexer.Add([]byte(`([Aa][Nn][Dd])`), token("AND_OPERATOR"))
 	lexer.Add([]byte(`\.\s*==\s*`), token("EQUALS_SELF_OPERATOR"))
 	lexer.Add([]byte(`\s*==\s*`), token("EQUALS_OPERATOR"))
+	lexer.Add([]byte(`\s*:=\s*`), token("ASSIGN_OPERATOR"))
 	lexer.Add([]byte(`\[-?[0-9]+\]`), numberToken("ARRAY_INDEX", true))
 	lexer.Add([]byte(`-?[0-9]+`), numberToken("ARRAY_INDEX", false))
 	lexer.Add([]byte("( |\t|\n|\r)+"), skip)