Variable loop wip

pkg/yqlib/lib.go

@@ -53,7 +53,7 @@ var subtractAssignOpType = &operationType{Type: "SUBTRACT_ASSIGN", NumArgs: 2, P
 
 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 assignVariableOpType = &operationType{Type: "ASSIGN_VARIABLE", NumArgs: 2, Precedence: 40, Handler: useWithPipe}
 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}

pkg/yqlib/operator_pipe.go

@@ -2,6 +2,10 @@ package yqlib
 
 func pipeOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
 
+	if expressionNode.LHS.Operation.OperationType == assignVariableOpType {
+		return variableLoop(d, context, expressionNode)
+	}
+
 	//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|...

pkg/yqlib/operator_variables.go

@@ -19,24 +19,54 @@ type assignVarPreferences struct {
 	IsReference bool
 }
 
-func assignVariableOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
+func useWithPipe(d *dataTreeNavigator, context Context, originalExp *ExpressionNode) (Context, error) {
-	lhs, err := d.GetMatchingNodes(context.ReadOnlyClone(), expressionNode.LHS)
+	return Context{}, fmt.Errorf("must use variable with a pipe, e.g. `exp as $x | ...`")
+}
+
+func variableLoop(d *dataTreeNavigator, context Context, originalExp *ExpressionNode) (Context, error) {
+	log.Debug("variable loop!")
+	variableExp := originalExp.LHS
+	lhs, err := d.GetMatchingNodes(context.ReadOnlyClone(), variableExp.LHS)
 	if err != nil {
 		return Context{}, err
 	}
-	if expressionNode.RHS.Operation.OperationType.Type != "GET_VARIABLE" {
+	if variableExp.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
+	variableName := variableExp.RHS.Operation.StringValue
 
-	prefs := expressionNode.Operation.Preferences.(assignVarPreferences)
+	prefs := variableExp.Operation.Preferences.(assignVarPreferences)
 
-	var variableValue *list.List
+	results := list.New()
-	if prefs.IsReference {
+
-		variableValue = lhs.MatchingNodes
+	// now we loop over lhs, set variable to each result and calculate originalExp.Rhs
-	} else {
+	for el := lhs.MatchingNodes.Front(); el != nil; el = el.Next() {
-		variableValue = lhs.DeepClone().MatchingNodes
+		var variableValue = list.New()
+		if prefs.IsReference {
+			variableValue.PushBack(el.Value)
+		} else {
+			copy, err := el.Value.(*CandidateNode).Copy()
+			if err != nil {
+				return Context{}, err
+			}
+			variableValue.PushBack(copy)
+		}
+		log.Debug("PROCESSING VARIABLE: ", NodeToString(el.Value.(*CandidateNode)))
+		newContext := context.ChildContext(context.MatchingNodes)
+		newContext.SetVariable(variableName, variableValue)
+
+		rhs, err := d.GetMatchingNodes(newContext, originalExp.RHS)
+		if err != nil {
+			return Context{}, err
+		}
+		results.PushBackList(rhs.MatchingNodes)
 	}
-	context.SetVariable(variableName, variableValue)
+
-	return context, nil
+	// if there is no LHS - then I guess we just calculate originalExp.Rhs
+	if lhs.MatchingNodes.Len() == 0 {
+		return d.GetMatchingNodes(context, originalExp.RHS)
+	}
+
+	return context.ChildContext(results), nil
+
 }

pkg/yqlib/operator_variables_test.go

@@ -8,7 +8,7 @@ var variableOperatorScenarios = []expressionScenario{
 	{
 		skipDoc:    true,
 		document:   `{}`,
-		expression: `.a.b as $foo`,
+		expression: `.a.b as $foo | .`,
 		expected: []string{
 			"D0, P[], (doc)::{}\n",
 		},
@@ -16,7 +16,7 @@ var variableOperatorScenarios = []expressionScenario{
 	{
 		document:      "a: [cat]",
 		skipDoc:       true,
-		expression:    "(.[] | {.name: .}) as $item",
+		expression:    "(.[] | {.name: .}) as $item | .",
 		expectedError: `cannot index array with 'name' (strconv.ParseInt: parsing "name": invalid syntax)`,
 	},
 	{
@@ -36,6 +36,14 @@ var variableOperatorScenarios = []expressionScenario{
 			"D0, P[1], (!!str)::dog\n",
 		},
 	},
+	{
+		skipDoc:    true,
+		document:   `[1, 2]`,
+		expression: `[.[] as $f | $f + 1]`,
+		expected: []string{
+			"D0, P[], (!!seq)::- 2\n- 3\n",
+		},
+	},
 	{
 		description:    "Using variables as a lookup",
 		subdescription: "Example taken from [jq](https://stedolan.github.io/jq/manual/#Variable/SymbolicBindingOperator:...as$identifier|...)",