diff --git a/pkg/yqlib/doc/operators/headers/pick.md b/pkg/yqlib/doc/operators/headers/pick.md new file mode 100644 index 00000000..f5a0d602 --- /dev/null +++ b/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. diff --git a/pkg/yqlib/doc/operators/pick.md b/pkg/yqlib/doc/operators/pick.md new file mode 100644 index 00000000..fd87a34b --- /dev/null +++ b/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 ` +{% 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 +``` + diff --git a/pkg/yqlib/expression_tokeniser.go b/pkg/yqlib/expression_tokeniser.go index ac8b7f8e..57c0bde9 100644 --- a/pkg/yqlib/expression_tokeniser.go +++ b/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})) diff --git a/pkg/yqlib/lib.go b/pkg/yqlib/lib.go index d57e5d4e..3c25aafd 100644 --- a/pkg/yqlib/lib.go +++ b/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, diff --git a/pkg/yqlib/operator_pick.go b/pkg/yqlib/operator_pick.go new file mode 100644 index 00000000..d4ae88ae --- /dev/null +++ b/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 +} diff --git a/pkg/yqlib/operator_pick_test.go b/pkg/yqlib/operator_pick_test.go new file mode 100644 index 00000000..f0e283fb --- /dev/null +++ b/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) +} diff --git a/pkg/yqlib/operators_test.go b/pkg/yqlib/operators_test.go index d3b8ad92..37e45c2f 100644 --- a/pkg/yqlib/operators_test.go +++ b/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" )