diff --git a/pkg/yqlib/lib.go b/pkg/yqlib/lib.go index 8b61eae3..65389263 100644 --- a/pkg/yqlib/lib.go +++ b/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} diff --git a/pkg/yqlib/operator_pipe.go b/pkg/yqlib/operator_pipe.go index 77813c67..beed4141 100644 --- a/pkg/yqlib/operator_pipe.go +++ b/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|... diff --git a/pkg/yqlib/operator_variables.go b/pkg/yqlib/operator_variables.go index 5550bded..3456666c 100644 --- a/pkg/yqlib/operator_variables.go +++ b/pkg/yqlib/operator_variables.go @@ -19,24 +19,54 @@ type assignVarPreferences struct { IsReference bool } -func assignVariableOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) { - lhs, err := d.GetMatchingNodes(context.ReadOnlyClone(), expressionNode.LHS) +func useWithPipe(d *dataTreeNavigator, context Context, originalExp *ExpressionNode) (Context, error) { + 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 - if prefs.IsReference { - variableValue = lhs.MatchingNodes - } else { - variableValue = lhs.DeepClone().MatchingNodes + results := list.New() + + // now we loop over lhs, set variable to each result and calculate originalExp.Rhs + for el := lhs.MatchingNodes.Front(); el != nil; el = el.Next() { + 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 + } diff --git a/pkg/yqlib/operator_variables_test.go b/pkg/yqlib/operator_variables_test.go index 7ab9ee30..b8fffdde 100644 --- a/pkg/yqlib/operator_variables_test.go +++ b/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|...)",