Added variables

pkg/yqlib/context.go

@@ -18,6 +18,20 @@ func (n *Context) SingleChildContext(candidate *CandidateNode) Context {
 	return n.ChildContext(list)
 }
 
+func (n *Context) GetVariable(name string) *list.List {
+	if n.Variables == nil {
+		return nil
+	}
+	return n.Variables[name]
+}
+
+func (n *Context) SetVariable(name string, value *list.List) {
+	if n.Variables == nil {
+		n.Variables = make(map[string]*list.List)
+	}
+	n.Variables[name] = value
+}
+
 func (n *Context) ChildContext(results *list.List) Context {
 	clone := Context{}
 	err := copier.Copy(&clone, n)

pkg/yqlib/expression_postfix.go

@@ -26,7 +26,7 @@ func popOpToResult(opStack []*token, result []*Operation) ([]*token, []*Operatio
 func (p *expressionPostFixerImpl) ConvertToPostfix(infixTokens []*token) ([]*Operation, error) {
 	var result []*Operation
 	// surround the whole thing with quotes
-	var opStack = []*token{&token{TokenType: openBracket}}
+	var opStack = []*token{{TokenType: openBracket}}
 	var tokens = append(infixTokens, &token{TokenType: closeBracket})
 
 	for _, currentToken := range tokens {

pkg/yqlib/expression_processing_test.go

@@ -17,35 +17,40 @@ var pathTests = []struct {
 		append(make([]interface{}, 0), "[", "]"),
 		append(make([]interface{}, 0), "EMPTY", "COLLECT", "SHORT_PIPE"),
 	},
+	{
+		`.b[.a]`,
+		append(make([]interface{}, 0), "b", "TRAVERSE_ARRAY", "[", "a", "]"),
+		append(make([]interface{}, 0), "b", "a", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY"),
+	},
 	{
 		`.[]`,
-		append(make([]interface{}, 0), "TRAVERSE_ARRAY", "[", "]"),
+		append(make([]interface{}, 0), "SELF", "TRAVERSE_ARRAY", "[", "]"),
-		append(make([]interface{}, 0), "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY"),
+		append(make([]interface{}, 0), "SELF", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY"),
 	},
 	{
 		`.a[]`,
-		append(make([]interface{}, 0), "a", "SHORT_PIPE", "TRAVERSE_ARRAY", "[", "]"),
+		append(make([]interface{}, 0), "a", "TRAVERSE_ARRAY", "[", "]"),
-		append(make([]interface{}, 0), "a", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "SHORT_PIPE"),
+		append(make([]interface{}, 0), "a", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY"),
 	},
 	{
 		`.a.[]`,
-		append(make([]interface{}, 0), "a", "SHORT_PIPE", "TRAVERSE_ARRAY", "[", "]"),
+		append(make([]interface{}, 0), "a", "TRAVERSE_ARRAY", "[", "]"),
-		append(make([]interface{}, 0), "a", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "SHORT_PIPE"),
+		append(make([]interface{}, 0), "a", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY"),
 	},
 	{
 		`.a[0]`,
-		append(make([]interface{}, 0), "a", "SHORT_PIPE", "TRAVERSE_ARRAY", "[", "0 (int64)", "]"),
+		append(make([]interface{}, 0), "a", "TRAVERSE_ARRAY", "[", "0 (int64)", "]"),
-		append(make([]interface{}, 0), "a", "0 (int64)", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "SHORT_PIPE"),
+		append(make([]interface{}, 0), "a", "0 (int64)", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY"),
 	},
 	{
 		`.a.[0]`,
-		append(make([]interface{}, 0), "a", "SHORT_PIPE", "TRAVERSE_ARRAY", "[", "0 (int64)", "]"),
+		append(make([]interface{}, 0), "a", "TRAVERSE_ARRAY", "[", "0 (int64)", "]"),
-		append(make([]interface{}, 0), "a", "0 (int64)", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "SHORT_PIPE"),
+		append(make([]interface{}, 0), "a", "0 (int64)", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY"),
 	},
 	{
 		`.a[].c`,
-		append(make([]interface{}, 0), "a", "SHORT_PIPE", "TRAVERSE_ARRAY", "[", "]", "SHORT_PIPE", "c"),
+		append(make([]interface{}, 0), "a", "TRAVERSE_ARRAY", "[", "]", "SHORT_PIPE", "c"),
-		append(make([]interface{}, 0), "a", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "SHORT_PIPE", "c", "SHORT_PIPE"),
+		append(make([]interface{}, 0), "a", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "c", "SHORT_PIPE"),
 	},
 	{
 		`[3]`,
@@ -69,18 +74,18 @@ var pathTests = []struct {
 	},
 	{
 		`.a | .[].b == "apple"`,
-		append(make([]interface{}, 0), "a", "PIPE", "TRAVERSE_ARRAY", "[", "]", "SHORT_PIPE", "b", "EQUALS", "apple (string)"),
+		append(make([]interface{}, 0), "a", "PIPE", "SELF", "TRAVERSE_ARRAY", "[", "]", "SHORT_PIPE", "b", "EQUALS", "apple (string)"),
-		append(make([]interface{}, 0), "a", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "b", "SHORT_PIPE", "apple (string)", "EQUALS", "PIPE"),
+		append(make([]interface{}, 0), "a", "SELF", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "b", "SHORT_PIPE", "apple (string)", "EQUALS", "PIPE"),
 	},
 	{
 		`(.a | .[].b) == "apple"`,
-		append(make([]interface{}, 0), "(", "a", "PIPE", "TRAVERSE_ARRAY", "[", "]", "SHORT_PIPE", "b", ")", "EQUALS", "apple (string)"),
+		append(make([]interface{}, 0), "(", "a", "PIPE", "SELF", "TRAVERSE_ARRAY", "[", "]", "SHORT_PIPE", "b", ")", "EQUALS", "apple (string)"),
-		append(make([]interface{}, 0), "a", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "b", "SHORT_PIPE", "PIPE", "apple (string)", "EQUALS"),
+		append(make([]interface{}, 0), "a", "SELF", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "b", "SHORT_PIPE", "PIPE", "apple (string)", "EQUALS"),
 	},
 	{
 		`.[] | select(. == "*at")`,
-		append(make([]interface{}, 0), "TRAVERSE_ARRAY", "[", "]", "PIPE", "SELECT", "(", "SELF", "EQUALS", "*at (string)", ")"),
+		append(make([]interface{}, 0), "SELF", "TRAVERSE_ARRAY", "[", "]", "PIPE", "SELECT", "(", "SELF", "EQUALS", "*at (string)", ")"),
-		append(make([]interface{}, 0), "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "SELF", "*at (string)", "EQUALS", "SELECT", "PIPE"),
+		append(make([]interface{}, 0), "SELF", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "SELF", "*at (string)", "EQUALS", "SELECT", "PIPE"),
 	},
 	{
 		`[true]`,
@@ -113,9 +118,9 @@ var pathTests = []struct {
 		append(make([]interface{}, 0), "a", "mike (string)", "CREATE_MAP", "COLLECT_OBJECT", "SHORT_PIPE"),
 	},
 	{
-		`{.a: .c, .b.[]: .f.g.[]}`,
+		`{.a: .c, .b.[]: .f.g[]}`,
-		append(make([]interface{}, 0), "{", "a", "CREATE_MAP", "c", "UNION", "b", "SHORT_PIPE", "TRAVERSE_ARRAY", "[", "]", "CREATE_MAP", "f", "SHORT_PIPE", "g", "SHORT_PIPE", "TRAVERSE_ARRAY", "[", "]", "}"),
+		append(make([]interface{}, 0), "{", "a", "CREATE_MAP", "c", "UNION", "b", "TRAVERSE_ARRAY", "[", "]", "CREATE_MAP", "f", "SHORT_PIPE", "g", "TRAVERSE_ARRAY", "[", "]", "}"),
-		append(make([]interface{}, 0), "a", "c", "CREATE_MAP", "b", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "SHORT_PIPE", "f", "g", "SHORT_PIPE", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "SHORT_PIPE", "CREATE_MAP", "UNION", "COLLECT_OBJECT", "SHORT_PIPE"),
+		append(make([]interface{}, 0), "a", "c", "CREATE_MAP", "b", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "f", "g", "EMPTY", "COLLECT", "SHORT_PIPE", "TRAVERSE_ARRAY", "SHORT_PIPE", "CREATE_MAP", "UNION", "COLLECT_OBJECT", "SHORT_PIPE"),
 	},
 	{
 		`explode(.a.b)`,

pkg/yqlib/expression_tokeniser.go

@@ -180,6 +180,19 @@ func stringValue(wrapped bool) lex.Action {
 	}
 }
 
+func getVariableOpToken() lex.Action {
+	return func(s *lex.Scanner, m *machines.Match) (interface{}, error) {
+		value := string(m.Bytes)
+
+		value = value[1 : len(value)-1]
+
+		getVarOperation := createValueOperation(value, value)
+		getVarOperation.OperationType = getVariableOpType
+
+		return &token{TokenType: operationToken, Operation: getVarOperation, CheckForPostTraverse: true}, nil
+	}
+}
+
 func envOp(strenv bool) lex.Action {
 	return func(s *lex.Scanner, m *machines.Match) (interface{}, error) {
 		value := string(m.Bytes)
@@ -304,6 +317,8 @@ func initLexer() (*lex.Lexer, error) {
 	lexer.Add([]byte(`\*[\+|\?]*`), multiplyWithPrefs())
 	lexer.Add([]byte(`\+`), opToken(addOpType))
 	lexer.Add([]byte(`\+=`), opToken(addAssignOpType))
+	lexer.Add([]byte(`\$[a-zA-Z_-0-9]+`), getVariableOpToken())
+	lexer.Add([]byte(`as`), opToken(assignVariableOpType))
 
 	err := lexer.Compile()
 	if err != nil {
@@ -369,6 +384,12 @@ func (p *expressionTokeniserImpl) handleToken(tokens []*token, index int, postPr
 		//need to put a traverse array then a collect currentToken
 		// do this by adding traverse then converting currentToken to collect
 
+		if index == 0 || tokens[index-1].TokenType != operationToken ||
+			tokens[index-1].Operation.OperationType != traversePathOpType {
+			op := &Operation{OperationType: selfReferenceOpType, StringValue: "SELF"}
+			postProcessedTokens = append(postProcessedTokens, &token{TokenType: operationToken, Operation: op})
+		}
+
 		op := &Operation{OperationType: traverseArrayOpType, StringValue: "TRAVERSE_ARRAY"}
 		postProcessedTokens = append(postProcessedTokens, &token{TokenType: operationToken, Operation: op})
 
@@ -395,18 +416,8 @@ func (p *expressionTokeniserImpl) handleToken(tokens []*token, index int, postPr
 	if index != len(tokens)-1 && currentToken.CheckForPostTraverse &&
 		tokens[index+1].TokenType == openCollect {
 
-		op := &Operation{OperationType: shortPipeOpType, Value: "PIPE"}
+		op := &Operation{OperationType: traverseArrayOpType}
 		postProcessedTokens = append(postProcessedTokens, &token{TokenType: operationToken, Operation: op})
-
-		op = &Operation{OperationType: traverseArrayOpType}
-		postProcessedTokens = append(postProcessedTokens, &token{TokenType: operationToken, Operation: op})
-	}
-	if index != len(tokens)-1 && currentToken.CheckForPostTraverse &&
-		tokens[index+1].TokenType == traverseArrayCollect {
-
-		op := &Operation{OperationType: shortPipeOpType, Value: "PIPE"}
-		postProcessedTokens = append(postProcessedTokens, &token{TokenType: operationToken, Operation: op})
-
 	}
 	return postProcessedTokens, skipNextToken
 }

pkg/yqlib/lib.go

@@ -35,6 +35,7 @@ var addAssignOpType = &operationType{Type: "ADD_ASSIGN", NumArgs: 2, Precedence:
 
 var assignAttributesOpType = &operationType{Type: "ASSIGN_ATTRIBUTES", NumArgs: 2, Precedence: 40, Handler: assignAttributesOperator}
 var assignStyleOpType = &operationType{Type: "ASSIGN_STYLE", NumArgs: 2, Precedence: 40, Handler: assignStyleOperator}
+var assignVariableOpType = &operationType{Type: "ASSIGN_VARIABLE", NumArgs: 2, Precedence: 40, Handler: assignVariableOperator}
 var assignTagOpType = &operationType{Type: "ASSIGN_TAG", NumArgs: 2, Precedence: 40, Handler: assignTagOperator}
 var assignCommentOpType = &operationType{Type: "ASSIGN_COMMENT", NumArgs: 2, Precedence: 40, Handler: assignCommentsOperator}
 var assignAnchorOpType = &operationType{Type: "ASSIGN_ANCHOR", NumArgs: 2, Precedence: 40, Handler: assignAnchorOperator}
@@ -52,6 +53,7 @@ var shortPipeOpType = &operationType{Type: "SHORT_PIPE", NumArgs: 2, Precedence:
 var lengthOpType = &operationType{Type: "LENGTH", NumArgs: 0, Precedence: 50, Handler: lengthOperator}
 var collectOpType = &operationType{Type: "COLLECT", NumArgs: 0, Precedence: 50, Handler: collectOperator}
 var splitDocumentOpType = &operationType{Type: "SPLIT_DOC", NumArgs: 0, Precedence: 50, Handler: splitDocumentOperator}
+var getVariableOpType = &operationType{Type: "GET_VARIABLE", NumArgs: 0, Precedence: 50, Handler: getVariableOperator}
 var getStyleOpType = &operationType{Type: "GET_STYLE", NumArgs: 0, Precedence: 50, Handler: getStyleOperator}
 var getTagOpType = &operationType{Type: "GET_TAG", NumArgs: 0, Precedence: 50, Handler: getTagOperator}
 var getCommentOpType = &operationType{Type: "GET_COMMENT", NumArgs: 0, Precedence: 50, Handler: getCommentsOperator}
@@ -71,7 +73,7 @@ var keysOpType = &operationType{Type: "KEYS", NumArgs: 0, Precedence: 50, Handle
 
 var collectObjectOpType = &operationType{Type: "COLLECT_OBJECT", NumArgs: 0, Precedence: 50, Handler: collectObjectOperator}
 var traversePathOpType = &operationType{Type: "TRAVERSE_PATH", NumArgs: 0, Precedence: 50, Handler: traversePathOperator}
-var traverseArrayOpType = &operationType{Type: "TRAVERSE_ARRAY", NumArgs: 1, Precedence: 50, Handler: traverseArrayOperator}
+var traverseArrayOpType = &operationType{Type: "TRAVERSE_ARRAY", NumArgs: 2, Precedence: 50, Handler: traverseArrayOperator}
 
 var selfReferenceOpType = &operationType{Type: "SELF", NumArgs: 0, Precedence: 50, Handler: selfOperator}
 var valueOpType = &operationType{Type: "VALUE", NumArgs: 0, Precedence: 50, Handler: valueOperator}

pkg/yqlib/operator_traverse_path.go

@@ -20,7 +20,7 @@ func splat(d *dataTreeNavigator, context Context, prefs traversePreferences) (Co
 }
 
 func traversePathOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
-	log.Debugf("-- Traversing")
+	log.Debugf("-- traversePathOperator")
 	var matches = list.New()
 
 	for el := context.MatchingNodes.Front(); el != nil; el = el.Next() {
@@ -75,6 +75,16 @@ func traverse(d *dataTreeNavigator, context Context, matchingNode *CandidateNode
 }
 
 func traverseArrayOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
+
+	//lhs may update the variable context, we should pass that into the RHS
+	// BUT we still return the original context back (see jq)
+	// https://stedolan.github.io/jq/manual/#Variable/SymbolicBindingOperator:...as$identifier|...
+
+	lhs, err := d.GetMatchingNodes(context, expressionNode.Lhs)
+	if err != nil {
+		return Context{}, err
+	}
+
 	// rhs is a collect expression that will yield indexes to retreive of the arrays
 
 	rhs, err := d.GetMatchingNodes(context, expressionNode.Rhs)
@@ -83,7 +93,13 @@ func traverseArrayOperator(d *dataTreeNavigator, context Context, expressionNode
 	}
 
 	var indicesToTraverse = rhs.MatchingNodes.Front().Value.(*CandidateNode).Node.Content
-	return traverseNodesWithArrayIndices(context, indicesToTraverse, traversePreferences{})
+
+	//now we traverse the result of the lhs against the indices we found
+	result, err := traverseNodesWithArrayIndices(lhs, indicesToTraverse, traversePreferences{})
+	if err != nil {
+		return Context{}, err
+	}
+	return context.ChildContext(result.MatchingNodes), nil
 }
 
 func traverseNodesWithArrayIndices(context Context, indicesToTraverse []*yaml.Node, prefs traversePreferences) (Context, error) {

pkg/yqlib/operator_traverse_path_test.go

@@ -63,6 +63,14 @@ var traversePathOperatorScenarios = []expressionScenario{
 			"D0, P[apple], (!!str)::crispy yum\n",
 		},
 	},
+	{
+		skipDoc:    true,
+		document:   `{b: apple, fruit: {apple: yum, banana: smooth}}`,
+		expression: `.fruit[.b]`,
+		expected: []string{
+			"D0, P[fruit apple], (!!str)::yum\n",
+		},
+	},
 	{
 		description:    "Children don't exist",
 		subdescription: "Nodes are added dynamically while traversing",
@@ -83,7 +91,7 @@ var traversePathOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{}`,
-		expression: `.a.[1]`,
+		expression: `.a[1]`,
 		expected: []string{
 			"D0, P[a 1], (!!null)::null\n",
 		},
@@ -154,7 +162,7 @@ var traversePathOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{a: &cat {c: frog}, b: *cat}`,
-		expression: `.b.[]`,
+		expression: `.b[]`,
 		expected: []string{
 			"D0, P[b c], (!!str)::frog\n",
 		},
@@ -236,7 +244,7 @@ var traversePathOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   mergeDocSample,
-		expression: `.foobar.[]`,
+		expression: `.foobar[]`,
 		expected: []string{
 			"D0, P[foobar c], (!!str)::foo_c\n",
 			"D0, P[foobar a], (!!str)::foo_a\n",
@@ -298,7 +306,7 @@ var traversePathOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   mergeDocSample,
-		expression: `.foobarList.[]`,
+		expression: `.foobarList[]`,
 		expected: []string{
 			"D0, P[foobarList b], (!!str)::bar_b\n",
 			"D0, P[foobarList a], (!!str)::foo_a\n",
@@ -344,7 +352,7 @@ var traversePathOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{a: [a,b,c]}`,
-		expression: `.a.[0, 2]`,
+		expression: `.a[0, 2]`,
 		expected: []string{
 			"D0, P[a 0], (!!str)::a\n",
 			"D0, P[a 2], (!!str)::c\n",
@@ -361,7 +369,7 @@ var traversePathOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{a: [a,b,c]}`,
-		expression: `.a.[-1]`,
+		expression: `.a[-1]`,
 		expected: []string{
 			"D0, P[a -1], (!!str)::c\n",
 		},
@@ -377,7 +385,7 @@ var traversePathOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{a: [a,b,c]}`,
-		expression: `.a.[-2]`,
+		expression: `.a[-2]`,
 		expected: []string{
 			"D0, P[a -2], (!!str)::b\n",
 		},
@@ -395,7 +403,7 @@ var traversePathOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{a: [a,b,c]}`,
-		expression: `.a.[]`,
+		expression: `.a[]`,
 		expected: []string{
 			"D0, P[a 0], (!!str)::a\n",
 			"D0, P[a 1], (!!str)::b\n",

pkg/yqlib/operator_variables.go

@@ -0,0 +1,29 @@
+package yqlib
+
+import (
+	"container/list"
+	"fmt"
+)
+
+func getVariableOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
+	variableName := expressionNode.Operation.StringValue
+	log.Debug("getVariableOperator %v", variableName)
+	result := context.GetVariable(variableName)
+	if result == nil {
+		result = list.New()
+	}
+	return context.ChildContext(result), nil
+}
+
+func assignVariableOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
+	lhs, err := d.GetMatchingNodes(context, expressionNode.Lhs)
+	if err != nil {
+		return Context{}, nil
+	}
+	if expressionNode.Rhs.Operation.OperationType.Type != "GET_VARIABLE" {
+		return Context{}, fmt.Errorf("RHS of 'as' operator must be a variable name e.g. $foo")
+	}
+	variableName := expressionNode.Rhs.Operation.StringValue
+	context.SetVariable(variableName, lhs.MatchingNodes)
+	return context, nil
+}

pkg/yqlib/operator_variables_test.go

@@ -0,0 +1,31 @@
+package yqlib
+
+import (
+	"testing"
+)
+
+var variableOperatorScenarios = []expressionScenario{
+	{
+		description: "Single value variable",
+		document:    `a: cat`,
+		expression:  `.a as $foo | $foo`,
+		expected: []string{
+			"D0, P[a], (!!str)::cat\n",
+		},
+	},
+	{
+		description: "Multi value variable",
+		document:    `[cat, dog]`,
+		expression:  `.[] as $foo | $foo`,
+		expected: []string{
+			"D0, P[0], (!!str)::cat\n",
+			"D0, P[1], (!!str)::dog\n",
+		},
+	},
+}
+
+func TestVariableOperatorScenarios(t *testing.T) {
+	for _, tt := range variableOperatorScenarios {
+		testScenario(t, &tt)
+	}
+}