Repository: mikefarah/yq
Pull request #1403 (2022-10-28 09:39:14)

Slice array


Patch:
diff --git a/examples/array.yaml b/examples/array.yaml
index 47f04423..07b8808e 100644
--- a/examples/array.yaml
+++ b/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]
\ No newline at end of file
diff --git a/pkg/yqlib/doc/operators/slice-array.md b/pkg/yqlib/doc/operators/slice-array.md
new file mode 100644
index 00000000..f0cdc261
--- /dev/null
+++ b/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
+```
+
diff --git a/pkg/yqlib/lexer.go b/pkg/yqlib/lexer.go
index d4ab3409..92c3720f 100644
--- a/pkg/yqlib/lexer.go
+++ b/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 {
diff --git a/pkg/yqlib/lexer_participle.go b/pkg/yqlib/lexer_participle.go
index ab73a475..89fe682e 100644
--- a/pkg/yqlib/lexer_participle.go
+++ b/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)
diff --git a/pkg/yqlib/lib.go b/pkg/yqlib/lib.go
index 461d15ea..76e52b03 100644
--- a/pkg/yqlib/lib.go
+++ b/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
diff --git a/pkg/yqlib/operator_slice.go b/pkg/yqlib/operator_slice.go
new file mode 100644
index 00000000..8982c49f
--- /dev/null
+++ b/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
+}
diff --git a/pkg/yqlib/operator_traverse_path.go b/pkg/yqlib/operator_traverse_path.go
index ca7141aa..67d19742 100644
--- a/pkg/yqlib/operator_traverse_path.go
+++ b/pkg/yqlib/operator_traverse_path.go
@@ -76,7 +76,6 @@ func traverse(context Context, matchingNode *CandidateNode, operation *Operation
 }
 
 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|...


Test patch:
diff --git a/pkg/yqlib/lexer_participle_test.go b/pkg/yqlib/lexer_participle_test.go
index 3a8739ed..59fc3b1e 100644
--- a/pkg/yqlib/lexer_participle_test.go
+++ b/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{
diff --git a/pkg/yqlib/operator_slice_test.go b/pkg/yqlib/operator_slice_test.go
new file mode 100644
index 00000000..dca58b24
--- /dev/null
+++ b/pkg/yqlib/operator_slice_test.go
@@ -0,0 +1,81 @@
+package yqlib
+
+import "testing"
+
+var sliceArrayScenarios = []expressionScenario{
+	{
+		description: "Slicing arrays",
+		document:    `[cat, dog, frog, cow]`,
+		expression:  `.[1:3]`,
+		expected: []string{
+			"D0, P[], (!!seq)::- dog\n- frog\n",
+		},
+	},
+	{
+		description:    "Slicing arrays - without the first number",
+		subdescription: "Starts from the start of the array",
+		document:       `[cat, dog, frog, cow]`,
+		expression:     `.[:2]`,
+		expected: []string{
+			"D0, P[], (!!seq)::- cat\n- dog\n",
+		},
+	},
+	{
+		description:    "Slicing arrays - without the second number",
+		subdescription: "Finishes at the end of the array",
+		document:       `[cat, dog, frog, cow]`,
+		expression:     `.[2:]`,
+		expected: []string{
+			"D0, P[], (!!seq)::- frog\n- cow\n",
+		},
+	},
+	{
+		description: "Slicing arrays - use negative numbers to count backwards from the end",
+		document:    `[cat, dog, frog, cow]`,
+		expression:  `.[1:-1]`,
+		expected: []string{
+			"D0, P[], (!!seq)::- dog\n- frog\n",
+		},
+	},
+	{
+		skipDoc:    true,
+		document:   `[[cat, dog, frog, cow], [apple, banana, grape, mango]]`,
+		expression: `.[] | .[1:3]`,
+		expected: []string{
+			"D0, P[0], (!!seq)::- dog\n- frog\n",
+			"D0, P[1], (!!seq)::- banana\n- grape\n",
+		},
+	},
+	{
+		skipDoc:    true,
+		document:   `[[cat, dog, frog, cow], [apple, banana, grape, mango]]`,
+		expression: `.[] | .[-2:-1]`,
+		expected: []string{
+			"D0, P[0], (!!seq)::- frog\n",
+			"D0, P[1], (!!seq)::- grape\n",
+		},
+	},
+	{
+		skipDoc:    true,
+		document:   `[cat1, cat2, cat3, cat4, cat5, cat6, cat7, cat8, cat9, cat10, cat11]`,
+		expression: `.[10:11]`,
+		expected: []string{
+			"D0, P[], (!!seq)::- cat11\n",
+		},
+	},
+	{
+		skipDoc:    true,
+		document:   `[cat1, cat2, cat3, cat4, cat5, cat6, cat7, cat8, cat9, cat10, cat11]`,
+		expression: `.[-11:-10]`,
+		expected: []string{
+			"D0, P[], (!!seq)::- cat1\n",
+		},
+	},
+}
+
+func TestSliceOperatorScenarios(t *testing.T) {
+	for _, tt := range sliceArrayScenarios {
+		testScenario(t, &tt)
+	}
+	documentOperatorScenarios(t, "slice-array", sliceArrayScenarios)
+}
diff --git a/pkg/yqlib/operators_test.go b/pkg/yqlib/operators_test.go
index 79757895..f3363dc6 100644
--- a/pkg/yqlib/operators_test.go
+++ b/pkg/yqlib/operators_test.go
@@ -31,7 +31,7 @@ type expressionScenario struct {
 }
 
 func TestMain(m *testing.M) {
-	logging.SetLevel(logging.ERROR, "")
+	logging.SetLevel(logging.DEBUG, "")
 	Now = func() time.Time {
 		return time.Date(2021, time.May, 19, 1, 2, 3, 4, time.UTC)
 	}