diff --git a/pkg/yqlib/data_tree_navigator_test.go b/pkg/yqlib/data_tree_navigator_test.go index 571a93e..d4b18b3 100644 --- a/pkg/yqlib/data_tree_navigator_test.go +++ b/pkg/yqlib/data_tree_navigator_test.go @@ -2,36 +2,10 @@ package yqlib import ( "container/list" - "strings" - "testing" - - yaml "gopkg.in/yaml.v3" ) -var treeNavigator = NewDataTreeNavigator(NavigationPrefs{}) var treeCreator = NewPathTreeCreator() -func readDoc(t *testing.T, content string) *list.List { - inputList := list.New() - if content == "" { - return inputList - } - decoder := yaml.NewDecoder(strings.NewReader(content)) - var dataBucket yaml.Node - err := decoder.Decode(&dataBucket) - if err != nil { - t.Error(content) - t.Error(err) - } - - inputList.PushBack(&CandidateNode{ - Document: 0, - Filename: "test.yml", - Node: &dataBucket, - }) - return inputList -} - func resultsToString(results *list.List) []string { var pretty []string = make([]string, 0) for el := results.Front(); el != nil; el = el.Next() { diff --git a/pkg/yqlib/document_index_operator.go b/pkg/yqlib/document_index_operator.go new file mode 100644 index 0000000..52ec7a5 --- /dev/null +++ b/pkg/yqlib/document_index_operator.go @@ -0,0 +1,20 @@ +package yqlib + +import ( + "container/list" + "fmt" + + "gopkg.in/yaml.v3" +) + +func GetDocumentIndexOperator(d *dataTreeNavigator, matchingNodes *list.List, pathNode *PathTreeNode) (*list.List, error) { + var results = list.New() + + for el := matchingNodes.Front(); el != nil; el = el.Next() { + candidate := el.Value.(*CandidateNode) + node := &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprintf("%v", candidate.Document), Tag: "!!int"} + scalar := &CandidateNode{Node: node, Document: candidate.Document, Path: candidate.Path} + results.PushBack(scalar) + } + return results, nil +} diff --git a/pkg/yqlib/document_index_operator_test.go b/pkg/yqlib/document_index_operator_test.go new file mode 100644 index 0000000..1fbea1b --- /dev/null +++ b/pkg/yqlib/document_index_operator_test.go @@ -0,0 +1,32 @@ +package yqlib + +import ( + "testing" +) + +var documentIndexScenarios = []expressionScenario{ + { + description: "Retrieve a document index", + document: "a: cat\n---\na: frog\n", + expression: `.a | documentIndex`, + expected: []string{ + "D0, P[a], (!!int)::0\n", + "D1, P[a], (!!int)::1\n", + }, + }, + { + description: "Filter by document index", + document: "a: cat\n---\na: frog\n", + expression: `select(. | documentIndex == 1)`, + expected: []string{ + "D1, P[], (doc)::a: frog\n", + }, + }, +} + +func TestDocumentIndexScenarios(t *testing.T) { + for _, tt := range documentIndexScenarios { + testScenario(t, &tt) + } + documentScenarios(t, "Document Index Operator", documentIndexScenarios) +} diff --git a/pkg/yqlib/lib.go b/pkg/yqlib/lib.go index fadaf50..6982af1 100644 --- a/pkg/yqlib/lib.go +++ b/pkg/yqlib/lib.go @@ -18,7 +18,6 @@ type OperationType struct { // operators TODO: // - generator doc from operator tests -// - set comments not recursive // - documentIndex - retrieves document index, can be used with select // - mergeAppend (merges and appends arrays) // - mergeEmpty (sets only if the document is empty, do I do that now?) @@ -48,6 +47,7 @@ var Length = &OperationType{Type: "LENGTH", NumArgs: 0, Precedence: 50, Handler: var Collect = &OperationType{Type: "COLLECT", NumArgs: 0, Precedence: 50, Handler: CollectOperator} var GetStyle = &OperationType{Type: "GET_STYLE", NumArgs: 0, Precedence: 50, Handler: GetStyleOperator} var GetComment = &OperationType{Type: "GET_COMMENT", NumArgs: 0, Precedence: 50, Handler: GetCommentsOperator} +var GetDocumentIndex = &OperationType{Type: "GET_DOCUMENT_INDEX", NumArgs: 0, Precedence: 50, Handler: GetDocumentIndexOperator} var Explode = &OperationType{Type: "EXPLODE", NumArgs: 1, Precedence: 50, Handler: ExplodeOperator} diff --git a/pkg/yqlib/operators_test.go b/pkg/yqlib/operators_test.go index 03a3259..692d4ea 100644 --- a/pkg/yqlib/operators_test.go +++ b/pkg/yqlib/operators_test.go @@ -5,6 +5,7 @@ import ( "bytes" "fmt" "os" + "strings" "testing" "github.com/mikefarah/yq/v4/test" @@ -19,14 +20,12 @@ type expressionScenario struct { } func testScenario(t *testing.T, s *expressionScenario) { - - nodes := readDoc(t, s.document) - path, errPath := treeCreator.ParsePath(s.expression) + node, errPath := treeCreator.ParsePath(s.expression) if errPath != nil { t.Error(errPath) return } - results, errNav := treeNavigator.GetMatchingNodes(nodes, path) + results, errNav := EvaluateStream("sample.yaml", strings.NewReader(s.document), node) if errNav != nil { t.Error(errNav) @@ -67,14 +66,13 @@ func documentScenarios(t *testing.T, title string, scenarios []expressionScenari w.WriteString(fmt.Sprintf("Result\n")) - nodes := readDoc(t, s.document) - path, errPath := treeCreator.ParsePath(s.expression) + node, errPath := treeCreator.ParsePath(s.expression) if errPath != nil { t.Error(errPath) return } var output bytes.Buffer - results, err := treeNavigator.GetMatchingNodes(nodes, path) + results, err := EvaluateStream("sample.yaml", strings.NewReader(s.document), node) printer.PrintResults(results, bufio.NewWriter(&output)) w.WriteString(fmt.Sprintf("```yaml\n%v```\n", output.String())) diff --git a/pkg/yqlib/path_tokeniser.go b/pkg/yqlib/path_tokeniser.go index 4fa41c0..a952ff6 100644 --- a/pkg/yqlib/path_tokeniser.go +++ b/pkg/yqlib/path_tokeniser.go @@ -191,6 +191,8 @@ func initLexer() (*lex.Lexer, error) { lexer.Add([]byte(`or`), opToken(Or)) lexer.Add([]byte(`not`), opToken(Not)) + lexer.Add([]byte(`documentIndex`), opToken(GetDocumentIndex)) + lexer.Add([]byte(`style\s*=`), opToken(AssignStyle)) lexer.Add([]byte(`style`), opToken(GetStyle)) diff --git a/pkg/yqlib/printer.go b/pkg/yqlib/printer.go index 55dcec6..9aabd7b 100644 --- a/pkg/yqlib/printer.go +++ b/pkg/yqlib/printer.go @@ -52,7 +52,7 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List, writer io.Writer return nil } - var previousDocIndex uint = 0 + var previousDocIndex uint = matchingNodes.Front().Value.(*CandidateNode).Document for el := matchingNodes.Front(); el != nil; el = el.Next() { mappedDoc := el.Value.(*CandidateNode) diff --git a/pkg/yqlib/utils.go b/pkg/yqlib/utils.go index 73408d4..459a918 100644 --- a/pkg/yqlib/utils.go +++ b/pkg/yqlib/utils.go @@ -10,7 +10,9 @@ import ( yaml "gopkg.in/yaml.v3" ) -func readStream(filename string) (*yaml.Decoder, error) { +var treeNavigator = NewDataTreeNavigator(NavigationPrefs{}) + +func readStream(filename string) (io.Reader, error) { if filename == "" { return nil, errors.New("Must provide filename") } @@ -26,22 +28,15 @@ func readStream(filename string) (*yaml.Decoder, error) { defer safelyCloseFile(file) stream = file } - return yaml.NewDecoder(stream), nil + return stream, nil } -// put this in lib -func Evaluate(filename string, node *PathTreeNode) (*list.List, error) { - - var treeNavigator = NewDataTreeNavigator(NavigationPrefs{}) - +func EvaluateStream(filename string, reader io.Reader, node *PathTreeNode) (*list.List, error) { var matchingNodes = list.New() var currentIndex uint = 0 - var decoder, err = readStream(filename) - if err != nil { - return nil, err - } + decoder := yaml.NewDecoder(reader) for { var dataBucket yaml.Node errorReading := decoder.Decode(&dataBucket) @@ -66,8 +61,16 @@ func Evaluate(filename string, node *PathTreeNode) (*list.List, error) { matchingNodes.PushBackList(newMatches) currentIndex = currentIndex + 1 } +} + +func Evaluate(filename string, node *PathTreeNode) (*list.List, error) { + + var reader, err = readStream(filename) + if err != nil { + return nil, err + } + return EvaluateStream(filename, reader, node) - return matchingNodes, nil } func safelyRenameFile(from string, to string) {