Added pick operator

pkg/yqlib/doc/operators/headers/pick.md

@@ -0,0 +1,5 @@
+# Pick
+
+Filter a map by the specified list of keys. Map is returned with the key in the order of the pick list.
+
+Similarly, you can filter a map by the specified list of indices.

pkg/yqlib/doc/operators/pick.md

@@ -0,0 +1,53 @@
+# Pick
+
+Filter a map by the specified list of keys. Map is returned with the key in the order of the pick list.
+
+Similarly, you can filter a map by the specified list of indices.
+
+{% hint style="warning" %}
+Note that versions prior to 4.18 require the 'eval/e' command to be specified. 
+
+`yq e <exp> <file>`
+{% endhint %}
+
+## Pick keys from map
+Note that the order of the keys matches the pick order and non existent keys are skipped.
+
+Given a sample.yml file of:
+```yaml
+myMap:
+  cat: meow
+  dog: bark
+  thing: hamster
+  hamster: squeek
+```
+then
+```bash
+yq '.myMap |= pick(["hamster", "cat", "goat"])' sample.yml
+```
+will output
+```yaml
+myMap:
+  hamster: squeek
+  cat: meow
+```
+
+## Pick indices from array
+Note that the order of the indexes matches the pick order and non existent indexes are skipped.
+
+Given a sample.yml file of:
+```yaml
+- cat
+- leopard
+- lion
+```
+then
+```bash
+yq 'pick([2, 0, 734, -5])' sample.yml
+```
+will output
+```yaml
+- lion
+- cat
+```
+

pkg/yqlib/expression_tokeniser.go

@@ -322,6 +322,7 @@ func initLexer() (*lex.Lexer, error) {
 
 	lexer.Add([]byte(`map`), opToken(mapOpType))
 	lexer.Add([]byte(`map_values`), opToken(mapValuesOpType))
+	lexer.Add([]byte(`pick`), opToken(pickOpType))
 
 	lexer.Add([]byte(`flatten\([0-9]+\)`), flattenWithDepth())
 	lexer.Add([]byte(`flatten`), opTokenWithPrefs(flattenOpType, nil, flattenPreferences{depth: -1}))

pkg/yqlib/lib.go

@@ -85,6 +85,7 @@ var columnOpType = &operationType{Type: "LINE", NumArgs: 0, Precedence: 50, Hand
 
 var collectOpType = &operationType{Type: "COLLECT", NumArgs: 1, Precedence: 50, Handler: collectOperator}
 var mapOpType = &operationType{Type: "MAP", NumArgs: 1, Precedence: 50, Handler: mapOperator}
+var pickOpType = &operationType{Type: "PICK", NumArgs: 1, Precedence: 50, Handler: pickOperator}
 var evalOpType = &operationType{Type: "EVAL", NumArgs: 1, Precedence: 50, Handler: evalOperator}
 var mapValuesOpType = &operationType{Type: "MAP_VALUES", NumArgs: 1, Precedence: 50, Handler: mapValuesOperator}
 
@@ -197,6 +198,16 @@ func findInArray(array *yaml.Node, item *yaml.Node) int {
 	return -1
 }
 
+func findKeyInMap(array *yaml.Node, item *yaml.Node) int {
+
+	for index := 0; index < len(array.Content); index = index + 2 {
+		if recursiveNodeEqual(array.Content[index], item) {
+			return index
+		}
+	}
+	return -1
+}
+
 func recurseNodeObjectEqual(lhs *yaml.Node, rhs *yaml.Node) bool {
 	if len(lhs.Content) != len(rhs.Content) {
 		return false
@@ -274,11 +285,21 @@ func deepCloneContent(content []*yaml.Node) []*yaml.Node {
 	return clonedContent
 }
 
+func deepCloneNoContent(node *yaml.Node) *yaml.Node {
+	return deepCloneWithOptions(node, false)
+}
 func deepClone(node *yaml.Node) *yaml.Node {
+	return deepCloneWithOptions(node, true)
+}
+
+func deepCloneWithOptions(node *yaml.Node, cloneContent bool) *yaml.Node {
 	if node == nil {
 		return nil
 	}
-	clonedContent := deepCloneContent(node.Content)
+	var clonedContent []*yaml.Node
+	if cloneContent {
+		clonedContent = deepCloneContent(node.Content)
+	}
 	return &yaml.Node{
 		Content:     clonedContent,
 		Kind:        node.Kind,

pkg/yqlib/operator_pick.go

@@ -0,0 +1,86 @@
+package yqlib
+
+import (
+	"container/list"
+	"fmt"
+
+	yaml "gopkg.in/yaml.v3"
+)
+
+func pickMap(original *yaml.Node, indices *yaml.Node) *yaml.Node {
+
+	filteredContent := make([]*yaml.Node, 0)
+	for index := 0; index < len(indices.Content); index = index + 1 {
+		keyToFind := indices.Content[index]
+
+		indexInMap := findKeyInMap(original, keyToFind)
+		if indexInMap > -1 {
+			clonedKey := deepClone(original.Content[indexInMap])
+			clonedValue := deepClone(original.Content[indexInMap+1])
+			filteredContent = append(filteredContent, clonedKey, clonedValue)
+		}
+	}
+
+	newNode := deepCloneNoContent(original)
+	newNode.Content = filteredContent
+
+	return newNode
+}
+
+func pickSequence(original *yaml.Node, indices *yaml.Node) (*yaml.Node, error) {
+
+	filteredContent := make([]*yaml.Node, 0)
+	for index := 0; index < len(indices.Content); index = index + 1 {
+		_, indexInArray, err := parseInt(indices.Content[index].Value)
+		if err != nil {
+			return nil, fmt.Errorf("cannot index array with %v", indices.Content[index].Value)
+		}
+
+		if int(indexInArray) > -1 && int(indexInArray) < len(original.Content) {
+			filteredContent = append(filteredContent, deepClone(original.Content[indexInArray]))
+		}
+	}
+
+	newNode := deepCloneNoContent(original)
+	newNode.Content = filteredContent
+
+	return newNode, nil
+}
+
+func pickOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) {
+	log.Debugf("Pick")
+
+	contextIndicesToPick, err := d.GetMatchingNodes(context, expressionNode.RHS)
+
+	if err != nil {
+		return Context{}, err
+	}
+	indicesToPick := &yaml.Node{}
+	if contextIndicesToPick.MatchingNodes.Len() > 0 {
+		indicesToPick = contextIndicesToPick.MatchingNodes.Front().Value.(*CandidateNode).Node
+	}
+
+	var results = list.New()
+
+	for el := context.MatchingNodes.Front(); el != nil; el = el.Next() {
+		candidate := el.Value.(*CandidateNode)
+		node := unwrapDoc(candidate.Node)
+
+		var replacement *yaml.Node
+		if node.Tag == "!!map" {
+			replacement = pickMap(node, indicesToPick)
+		} else if node.Tag == "!!seq" {
+			replacement, err = pickSequence(node, indicesToPick)
+			if err != nil {
+				return Context{}, err
+			}
+
+		} else {
+			return Context{}, fmt.Errorf("cannot pick indicies from type %v (%v)", node.Tag, candidate.GetNicePath())
+		}
+
+		results.PushBack(candidate.CreateReplacement(replacement))
+	}
+
+	return context.ChildContext(results), nil
+}

pkg/yqlib/operator_pick_test.go

@@ -0,0 +1,33 @@
+package yqlib
+
+import (
+	"testing"
+)
+
+var pickOperatorScenarios = []expressionScenario{
+	{
+		description:    "Pick keys from map",
+		subdescription: "Note that the order of the keys matches the pick order and non existent keys are skipped.",
+		document:       "myMap: {cat: meow, dog: bark, thing: hamster, hamster: squeek}\n",
+		expression:     `.myMap |= pick(["hamster", "cat", "goat"])`,
+		expected: []string{
+			"D0, P[], (doc)::myMap: {hamster: squeek, cat: meow}\n",
+		},
+	},
+	{
+		description:    "Pick indices from array",
+		subdescription: "Note that the order of the indexes matches the pick order and non existent indexes are skipped.",
+		document:       `[cat, leopard, lion]`,
+		expression:     `pick([2, 0, 734, -5])`,
+		expected: []string{
+			"D0, P[], (!!seq)::[lion, cat]\n",
+		},
+	},
+}
+
+func TestPickOperatorScenarios(t *testing.T) {
+	for _, tt := range pickOperatorScenarios {
+		testScenario(t, &tt)
+	}
+	documentOperatorScenarios(t, "pick", pickOperatorScenarios)
+}

pkg/yqlib/operators_test.go

@@ -13,7 +13,7 @@ import (
 	"time"
 
 	"github.com/mikefarah/yq/v4/test"
-	"gopkg.in/op/go-logging.v1"
+	logging "gopkg.in/op/go-logging.v1"
 	yaml "gopkg.in/yaml.v3"
 )