Can use expressions in slice #1419

pkg/yqlib/doc/operators/slice-array.md

@@ -80,3 +80,26 @@ will output
 - frog
 ```
 
+## Inserting into the middle of an array
+using an expression to find the index
+
+Given a sample.yml file of:
+```yaml
+- cat
+- dog
+- frog
+- cow
+```
+then
+```bash
+yq '(.[] | select(. == "dog") | key + 1) as $pos | .[0:($pos)] + ["rabbit"] + .[$pos:]' sample.yml
+```
+will output
+```yaml
+- cat
+- dog
+- rabbit
+- frog
+- cow
+```
+

pkg/yqlib/lexer.go

@@ -97,6 +97,10 @@ func postProcessTokens(tokens []*token) []*token {
 	return postProcessedTokens
 }
 
+func tokenIsOpType(token *token, opType *operationType) bool {
+	return token.TokenType == operationToken && token.Operation.OperationType == opType
+}
+
 func handleToken(tokens []*token, index int, postProcessedTokens []*token) (tokensAccum []*token, skipNextToken bool) {
 	skipNextToken = false
 	currentToken := tokens[index]
@@ -120,9 +124,18 @@ func handleToken(tokens []*token, index int, postProcessedTokens []*token) (toke
 
 	}
 
+	if tokenIsOpType(currentToken, createMapOpType) {
+		log.Debugf("tokenIsOpType: createMapOpType")
+		// check the previous token is '[', means we are slice, but dont have a first number
+		if tokens[index-1].TokenType == traverseArrayCollect {
+			log.Debugf("previous token is : traverseArrayOpType")
+			// need to put the number 0 before this token, as that is implied
+			postProcessedTokens = append(postProcessedTokens, &token{TokenType: operationToken, Operation: createValueOperation(0, "0")})
+		}
+	}
+
 	if index != len(tokens)-1 && currentToken.AssignOperation != nil &&
-		tokens[index+1].TokenType == operationToken &&
+		tokenIsOpType(tokens[index+1], assignOpType) {
-		tokens[index+1].Operation.OperationType == assignOpType {
 		log.Debug("  its an update assign")
 		currentToken.Operation = currentToken.AssignOperation
 		currentToken.Operation.UpdateAssign = tokens[index+1].Operation.UpdateAssign
@@ -132,6 +145,17 @@ func handleToken(tokens []*token, index int, postProcessedTokens []*token) (toke
 	log.Debug("  adding token to the fixed list")
 	postProcessedTokens = append(postProcessedTokens, currentToken)
 
+	if tokenIsOpType(currentToken, createMapOpType) {
+		log.Debugf("tokenIsOpType: createMapOpType")
+		// check the next token is ']', means we are slice, but dont have a second number
+		if index != len(tokens)-1 && tokens[index+1].TokenType == closeCollect {
+			log.Debugf("next token is : closeCollect")
+			// need to put the number 0 before this token, as that is implied
+			lengthOp := &Operation{OperationType: lengthOpType}
+			postProcessedTokens = append(postProcessedTokens, &token{TokenType: operationToken, Operation: lengthOp})
+		}
+	}
+
 	if index != len(tokens)-1 &&
 		((currentToken.TokenType == openCollect && tokens[index+1].TokenType == closeCollect) ||
 			(currentToken.TokenType == openCollectObject && tokens[index+1].TokenType == closeCollectObject)) {
@@ -141,7 +165,8 @@ func handleToken(tokens []*token, index int, postProcessedTokens []*token) (toke
 	}
 
 	if index != len(tokens)-1 && currentToken.CheckForPostTraverse &&
-		((tokens[index+1].TokenType == operationToken && (tokens[index+1].Operation.OperationType == traversePathOpType)) ||
+
+		(tokenIsOpType(tokens[index+1], traversePathOpType) ||
 			(tokens[index+1].TokenType == traverseArrayCollect)) {
 		log.Debug("  adding pipe because the next thing is traverse")
 		op := &Operation{OperationType: shortPipeOpType, Value: "PIPE", StringValue: "."}

pkg/yqlib/lexer_participle.go

@@ -1,7 +1,6 @@
 package yqlib
 
 import (
-	"regexp"
 	"strconv"
 	"strings"
 
@@ -13,10 +12,6 @@ var participleYqRules = []*participleYqRule{
 	{"HEAD_COMMENT", `head_?comment|headComment`, opTokenWithPrefs(getCommentOpType, assignCommentOpType, commentOpPreferences{HeadComment: true}), 0},
 	{"FOOT_COMMENT", `foot_?comment|footComment`, opTokenWithPrefs(getCommentOpType, assignCommentOpType, commentOpPreferences{FootComment: true}), 0},
 
-	{"SliceArray", `\.\[-?[0-9]+:-?[0-9]+\]`, sliceArrayTwoNumbers(), 0},
-	{"SliceArraySecond", `\.\[\:-?[0-9]+\]`, sliceArraySecondNumberOnly(), 0},
-	{"SliceArrayFirst", `\.\[-?[0-9]+\:\]`, sliceArrayFirstNumberOnly(), 0},
-
 	{"OpenBracket", `\(`, literalToken(openBracket, false), 0},
 	{"CloseBracket", `\)`, literalToken(closeBracket, true), 0},
 	{"OpenTraverseArrayCollect", `\.\[`, literalToken(traverseArrayCollect, false), 0},
@@ -315,84 +310,6 @@ func flattenWithDepth() yqAction {
 	}
 }
 
-func sliceArrayTwoNumbers() yqAction {
-	return func(rawToken lexer.Token) (*token, error) {
-		value := rawToken.Value
-		sliceArrayNumbers := regexp.MustCompile(`\.\[(-?[0-9]+)\:(-?[0-9]+)\]`)
-		matches := sliceArrayNumbers.FindStringSubmatch(value)
-		log.Debug("sliceArrayTwoNumbers value: %v", value)
-		log.Debug("Matches: %v", matches)
-
-		firstNumber, err := parseInt(matches[1])
-		if err != nil {
-			return nil, err
-		}
-		secondNumber, err := parseInt(matches[2])
-		if err != nil {
-			return nil, err
-		}
-
-		prefs := sliceArrayPreferences{
-			firstNumber:         firstNumber,
-			secondNumber:        secondNumber,
-			secondNumberDefined: true,
-		}
-		log.Debug("%v", prefs)
-
-		op := &Operation{OperationType: sliceArrayOpType, Value: sliceArrayOpType.Type, StringValue: value, Preferences: prefs}
-		return &token{TokenType: operationToken, Operation: op}, nil
-	}
-}
-
-func sliceArraySecondNumberOnly() yqAction {
-	return func(rawToken lexer.Token) (*token, error) {
-		value := rawToken.Value
-		sliceArrayNumbers := regexp.MustCompile(`\.\[\:(-?[0-9]+)\]`)
-		matches := sliceArrayNumbers.FindStringSubmatch(value)
-		log.Debug("sliceArraySecondNumberOnly value: %v", value)
-		log.Debug("Matches: %v", matches)
-
-		secondNumber, err := parseInt(matches[1])
-		if err != nil {
-			return nil, err
-		}
-
-		prefs := sliceArrayPreferences{
-			firstNumber:         0,
-			secondNumber:        secondNumber,
-			secondNumberDefined: true,
-		}
-		log.Debug("%v", prefs)
-
-		op := &Operation{OperationType: sliceArrayOpType, Value: sliceArrayOpType.Type, StringValue: value, Preferences: prefs}
-		return &token{TokenType: operationToken, Operation: op}, nil
-	}
-}
-
-func sliceArrayFirstNumberOnly() yqAction {
-	return func(rawToken lexer.Token) (*token, error) {
-		value := rawToken.Value
-		sliceArrayNumbers := regexp.MustCompile(`\.\[(-?[0-9]+)\:\]`)
-		matches := sliceArrayNumbers.FindStringSubmatch(value)
-		log.Debug("sliceArrayFirstNumberOnly value: %v", value)
-		log.Debug("Matches: %v", matches)
-
-		firstNumber, err := parseInt(matches[1])
-		if err != nil {
-			return nil, err
-		}
-
-		prefs := sliceArrayPreferences{
-			firstNumber:         firstNumber,
-			secondNumberDefined: false,
-		}
-		log.Debug("%v", prefs)
-
-		op := &Operation{OperationType: sliceArrayOpType, Value: sliceArrayOpType.Type, StringValue: value, Preferences: prefs}
-		return &token{TokenType: operationToken, Operation: op}, nil
-	}
-}
-
 func assignAllCommentsOp(updateAssign bool) yqAction {
 	return func(rawToken lexer.Token) (*token, error) {
 		log.Debug("assignAllCommentsOp %v", rawToken.Value)

pkg/yqlib/lexer_participle_test.go

@@ -5,7 +5,7 @@ import (
 
 	"github.com/alecthomas/repr"
 	"github.com/mikefarah/yq/v4/test"
-	"gopkg.in/yaml.v3"
+	yaml "gopkg.in/yaml.v3"
 )
 
 type participleLexerScenario struct {
@@ -14,60 +14,125 @@ type participleLexerScenario struct {
 }
 
 var participleLexerScenarios = []participleLexerScenario{
-	{
-		expression: ".[1:3]",
-		tokens: []*token{
-			{
-				TokenType: operationToken,
-				Operation: &Operation{
-					OperationType: sliceArrayOpType,
-					Value:         "SLICE",
-					StringValue:   ".[1:3]",
-					Preferences:   sliceArrayPreferences{firstNumber: 1, secondNumber: 3, secondNumberDefined: true},
-				},
-			},
-		},
-	},
 	{
 		expression: ".[:3]",
 		tokens: []*token{
 			{
 				TokenType: operationToken,
 				Operation: &Operation{
-					OperationType: sliceArrayOpType,
+					OperationType: selfReferenceOpType,
-					Value:         "SLICE",
+					StringValue:   "SELF",
-					StringValue:   ".[:3]",
-					Preferences:   sliceArrayPreferences{firstNumber: 0, secondNumber: 3, secondNumberDefined: true},
 				},
 			},
+			{
+				TokenType: operationToken,
+				Operation: &Operation{
+					OperationType: traverseArrayOpType,
+					StringValue:   "TRAVERSE_ARRAY",
+				},
+			},
+			{
+				TokenType: openCollect,
+			},
+			{
+				TokenType: operationToken,
+				Operation: &Operation{
+					OperationType: valueOpType,
+					Value:         0,
+					StringValue:   "0",
+					CandidateNode: &CandidateNode{
+						Node: &yaml.Node{
+							Kind:  yaml.ScalarNode,
+							Tag:   "!!int",
+							Value: "0",
+						},
+					},
+				},
+			},
+			{
+				TokenType: operationToken,
+				Operation: &Operation{
+					OperationType: createMapOpType,
+					Value:         "CREATE_MAP",
+					StringValue:   ":",
+				},
+			},
+			{
+				TokenType: operationToken,
+				Operation: &Operation{
+					OperationType: valueOpType,
+					Value:         3,
+					StringValue:   "3",
+					CandidateNode: &CandidateNode{
+						Node: &yaml.Node{
+							Kind:  yaml.Kind(8),
+							Tag:   "!!int",
+							Value: "3",
+						},
+					},
+				},
+			},
+			{
+				TokenType:            closeCollect,
+				CheckForPostTraverse: true,
+				Match:                "]",
+			},
 		},
 	},
 	{
-		expression: ".[1:]",
+		expression: ".[-2:]",
 		tokens: []*token{
 			{
 				TokenType: operationToken,
 				Operation: &Operation{
-					OperationType: sliceArrayOpType,
+					OperationType: selfReferenceOpType,
-					Value:         "SLICE",
+					StringValue:   "SELF",
-					StringValue:   ".[1:]",
-					Preferences:   sliceArrayPreferences{firstNumber: 1, secondNumber: 0, secondNumberDefined: false},
 				},
 			},
-		},
-	},
-	{
-		expression: ".[-100:-54]",
-		tokens: []*token{
 			{
 				TokenType: operationToken,
 				Operation: &Operation{
-					OperationType: sliceArrayOpType,
+					OperationType: traverseArrayOpType,
-					Value:         "SLICE",
+					StringValue:   "TRAVERSE_ARRAY",
-					StringValue:   ".[-100:-54]",
-					Preferences:   sliceArrayPreferences{firstNumber: -100, secondNumber: -54, secondNumberDefined: true},
 				},
 			},
+			{
+				TokenType: openCollect,
+			},
+			{
+				TokenType: operationToken,
+				Operation: &Operation{
+					OperationType: valueOpType,
+					Value:         -2,
+					StringValue:   "-2",
+					CandidateNode: &CandidateNode{
+						Node: &yaml.Node{
+							Kind:  yaml.ScalarNode,
+							Tag:   "!!int",
+							Value: "-2",
+						},
+					},
+				},
+			},
+			{
+				TokenType: operationToken,
+				Operation: &Operation{
+					OperationType: createMapOpType,
+					Value:         "CREATE_MAP",
+					StringValue:   ":",
+				},
+			},
+			{
+				TokenType: operationToken,
+				Operation: &Operation{
+					OperationType: lengthOpType,
+				},
+			},
+			{
+				TokenType:            closeCollect,
+				CheckForPostTraverse: true,
+				Match:                "]",
+			},
 		},
 	},
 	{

pkg/yqlib/lib.go

@@ -83,7 +83,6 @@ var columnOpType = &operationType{Type: "LINE", NumArgs: 0, Precedence: 50, Hand
 var expressionOpType = &operationType{Type: "EXP", NumArgs: 0, Precedence: 50, Handler: expressionOperator}
 
 var collectOpType = &operationType{Type: "COLLECT", NumArgs: 1, Precedence: 50, Handler: collectOperator}
-var sliceArrayOpType = &operationType{Type: "SLICE", NumArgs: 0, Precedence: 50, Handler: sliceArrayOperator}
 var mapOpType = &operationType{Type: "MAP", NumArgs: 1, Precedence: 50, Handler: mapOperator}
 var errorOpType = &operationType{Type: "ERROR", NumArgs: 1, Precedence: 50, Handler: errorOperator}
 var pickOpType = &operationType{Type: "PICK", NumArgs: 1, Precedence: 50, Handler: pickOperator}

pkg/yqlib/operator_slice.go

@@ -2,43 +2,52 @@ package yqlib
 
 import (
 	"container/list"
+	"fmt"
 
 	yaml "gopkg.in/yaml.v3"
 )
 
-type sliceArrayPreferences struct {
+func getSliceNumber(d *dataTreeNavigator, context Context, node *CandidateNode, expressionNode *ExpressionNode) (int, error) {
-	firstNumber         int
+	result, err := d.GetMatchingNodes(context.SingleChildContext(node), expressionNode)
-	secondNumber        int
+	if err != nil {
-	secondNumberDefined bool
+		return 0, err
+	}
+	if result.MatchingNodes.Len() != 1 {
+		return 0, fmt.Errorf("expected to find 1 number, got %v instead", result.MatchingNodes.Len())
+	}
+	return parseInt(result.MatchingNodes.Front().Value.(*CandidateNode).Node.Value)
 }
 
 func sliceArrayOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
 
-	lhs, err := d.GetMatchingNodes(context, expressionNode.LHS)
+	log.Debug("slice array operator!")
-	if err != nil {
+	log.Debug("lhs: %v", expressionNode.LHS.Operation.toString())
-		return Context{}, err
+	log.Debug("rhs: %v", expressionNode.RHS.Operation.toString())
-	}
-	prefs := expressionNode.Operation.Preferences.(sliceArrayPreferences)
-	firstNumber := prefs.firstNumber
-	secondNumber := prefs.secondNumber
 
 	results := list.New()
 
-	for el := lhs.MatchingNodes.Front(); el != nil; el = el.Next() {
+	for el := context.MatchingNodes.Front(); el != nil; el = el.Next() {
 		lhsNode := el.Value.(*CandidateNode)
 		original := unwrapDoc(lhsNode.Node)
 
+		firstNumber, err := getSliceNumber(d, context, lhsNode, expressionNode.LHS)
+
+		if err != nil {
+			return Context{}, err
+		}
 		relativeFirstNumber := firstNumber
 		if relativeFirstNumber < 0 {
 			relativeFirstNumber = len(original.Content) + firstNumber
 		}
 
-		relativeSecondNumber := len(original.Content)
+		secondNumber, err := getSliceNumber(d, context, lhsNode, expressionNode.RHS)
-		if prefs.secondNumberDefined {
+		if err != nil {
-			relativeSecondNumber = secondNumber
+			return Context{}, err
-			if relativeSecondNumber < 0 {
+		}
-				relativeSecondNumber = len(original.Content) + secondNumber
+
-			}
+		relativeSecondNumber := secondNumber
+		if relativeSecondNumber < 0 {
+			relativeSecondNumber = len(original.Content) + secondNumber
 		}
 
 		log.Debug("calculateIndicesToTraverse: slice from %v to %v", relativeFirstNumber, relativeSecondNumber)

pkg/yqlib/operator_slice_test.go

@@ -37,6 +37,15 @@ var sliceArrayScenarios = []expressionScenario{
 			"D0, P[], (!!seq)::- dog\n- frog\n",
 		},
 	},
+	{
+		description:    "Inserting into the middle of an array",
+		subdescription: "using an expression to find the index",
+		document:       `[cat, dog, frog, cow]`,
+		expression:     `(.[] | select(. == "dog") | key + 1) as $pos | .[0:($pos)] + ["rabbit"] + .[$pos:]`,
+		expected: []string{
+			"D0, P[], (!!seq)::- cat\n- dog\n- rabbit\n- frog\n- cow\n",
+		},
+	},
 	{
 		skipDoc:    true,
 		document:   `[[cat, dog, frog, cow], [apple, banana, grape, mango]]`,

pkg/yqlib/operator_traverse_path.go

@@ -80,6 +80,10 @@ func traverseArrayOperator(d *dataTreeNavigator, context Context, expressionNode
 	// BUT we still return the original context back (see jq)
 	// https://stedolan.github.io/jq/manual/#Variable/SymbolicBindingOperator:...as$identifier|...
 
+	if expressionNode.RHS != nil && expressionNode.RHS.RHS != nil && expressionNode.RHS.RHS.Operation.OperationType == createMapOpType {
+		return sliceArrayOperator(d, context, expressionNode.RHS.RHS)
+	}
+
 	lhs, err := d.GetMatchingNodes(context, expressionNode.LHS)
 	if err != nil {
 		return Context{}, err