Slice array (#1403)

Author: mikefarah

examples/array.yaml

@@ -1,11 +1,2 @@
----
-- become: true
-  gather_facts: false
-  hosts: lalaland
-  name: "Apply smth"
-  roles:
-    - lala
-    - land
-  serial: 1
-- become: false
-  gather_facts: true
+- [cat, dog, frog, cow]
+- [apple, banana, grape, mango]

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

@@ -0,0 +1,77 @@
+
+## Slicing arrays
+Given a sample.yml file of:
+```yaml
+- cat
+- dog
+- frog
+- cow
+```
+then
+```bash
+yq '.[1:3]' sample.yml
+```
+will output
+```yaml
+- dog
+- frog
+```
+
+## Slicing arrays - without the first number
+Starts from the start of the array
+
+Given a sample.yml file of:
+```yaml
+- cat
+- dog
+- frog
+- cow
+```
+then
+```bash
+yq '.[:2]' sample.yml
+```
+will output
+```yaml
+- cat
+- dog
+```
+
+## Slicing arrays - without the second number
+Finishes at the end of the array
+
+Given a sample.yml file of:
+```yaml
+- cat
+- dog
+- frog
+- cow
+```
+then
+```bash
+yq '.[2:]' sample.yml
+```
+will output
+```yaml
+- frog
+- cow
+```
+
+## Slicing arrays - use negative numbers to count backwards from the end
+Given a sample.yml file of:
+```yaml
+- cat
+- dog
+- frog
+- cow
+```
+then
+```bash
+yq '.[1:-1]' sample.yml
+```
+will output
+```yaml
+- dog
+- frog
+```
+

pkg/yqlib/lexer.go

@@ -3,7 +3,6 @@ package yqlib
 import (
 	"fmt"
 	"regexp"
-	"strconv"
 )
 
 type expressionTokeniser interface {
@@ -64,11 +63,11 @@ func unwrap(value string) string {
 func extractNumberParameter(value string) (int, error) {
 	parameterParser := regexp.MustCompile(`.*\(([0-9]+)\)`)
 	matches := parameterParser.FindStringSubmatch(value)
-	var indent, errParsingInt = strconv.ParseInt(matches[1], 10, 32)
+	var indent, errParsingInt = parseInt(matches[1])
 	if errParsingInt != nil {
 		return 0, errParsingInt
 	}
-	return int(indent), nil
+	return indent, nil
 }
 
 func hasOptionParameter(value string, option string) bool {

pkg/yqlib/lexer_participle.go

@@ -1,6 +1,7 @@
 package yqlib
 
 import (
+	"regexp"
 	"strconv"
 	"strings"
 
@@ -12,6 +13,10 @@ 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},
@@ -300,6 +305,84 @@ 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

@@ -14,6 +14,62 @@ 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,
+					Value:         "SLICE",
+					StringValue:   ".[:3]",
+					Preferences:   sliceArrayPreferences{firstNumber: 0, secondNumber: 3, secondNumberDefined: true},
+				},
+			},
+		},
+	},
+	{
+		expression: ".[1:]",
+		tokens: []*token{
+			{
+				TokenType: operationToken,
+				Operation: &Operation{
+					OperationType: sliceArrayOpType,
+					Value:         "SLICE",
+					StringValue:   ".[1:]",
+					Preferences:   sliceArrayPreferences{firstNumber: 1, secondNumber: 0, secondNumberDefined: false},
+				},
+			},
+		},
+	},
+	{
+		expression: ".[-100:-54]",
+		tokens: []*token{
+			{
+				TokenType: operationToken,
+				Operation: &Operation{
+					OperationType: sliceArrayOpType,
+					Value:         "SLICE",
+					StringValue:   ".[-100:-54]",
+					Preferences:   sliceArrayPreferences{firstNumber: -100, secondNumber: -54, secondNumberDefined: true},
+				},
+			},
+		},
+	},
 	{
 		expression: ".a",
 		tokens: []*token{

pkg/yqlib/lib.go

@@ -81,6 +81,7 @@ var lineOpType = &operationType{Type: "LINE", NumArgs: 0, Precedence: 50, Handle
 var columnOpType = &operationType{Type: "LINE", NumArgs: 0, Precedence: 50, Handler: columnOperator}
 
 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}
@@ -352,8 +353,8 @@ func parseInt(numberString string) (int, error) {
 
 	if err != nil {
 		return 0, err
-	} else if parsed > math.MaxInt {
-		return 0, fmt.Errorf("%v is too big (larger than %v)", parsed, math.MaxInt)
+	} else if parsed > math.MaxInt || parsed < math.MinInt {
+		return 0, fmt.Errorf("%v is not within [%v, %v]", parsed, math.MinInt, math.MaxInt)
 	}
 
 	return int(parsed), err

pkg/yqlib/operator_slice.go

@@ -0,0 +1,63 @@
+package yqlib
+
+import (
+	"container/list"
+
+	yaml "gopkg.in/yaml.v3"
+)
+
+type sliceArrayPreferences struct {
+	firstNumber         int
+	secondNumber        int
+	secondNumberDefined bool
+}
+
+func sliceArrayOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
+
+	lhs, err := d.GetMatchingNodes(context, expressionNode.LHS)
+	if err != nil {
+		return Context{}, err
+	}
+	prefs := expressionNode.Operation.Preferences.(sliceArrayPreferences)
+	firstNumber := prefs.firstNumber
+	secondNumber := prefs.secondNumber
+
+	results := list.New()
+
+	for el := lhs.MatchingNodes.Front(); el != nil; el = el.Next() {
+		lhsNode := el.Value.(*CandidateNode)
+		original := unwrapDoc(lhsNode.Node)
+
+		relativeFirstNumber := firstNumber
+		if relativeFirstNumber < 0 {
+			relativeFirstNumber = len(original.Content) + firstNumber
+		}
+
+		relativeSecondNumber := len(original.Content)
+		if prefs.secondNumberDefined {
+			relativeSecondNumber = secondNumber
+			if relativeSecondNumber < 0 {
+				relativeSecondNumber = len(original.Content) + secondNumber
+			}
+		}
+
+		log.Debug("calculateIndicesToTraverse: slice from %v to %v", relativeFirstNumber, relativeSecondNumber)
+
+		var newResults []*yaml.Node
+		for i := relativeFirstNumber; i < relativeSecondNumber; i++ {
+			newResults = append(newResults, original.Content[i])
+		}
+
+		slicedArrayNode := &yaml.Node{
+			Kind:    yaml.SequenceNode,
+			Tag:     original.Tag,
+			Content: newResults,
+		}
+		results.PushBack(lhsNode.CreateReplacement(slicedArrayNode))
+
+	}
+
+	// result is now the context that has the nodes we need to put back into a sequence.
+	//what about multiple arrays in the context? I think we need to create an array for each one
+	return context.ChildContext(results), nil
+}