document index

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() {

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
+}

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)
+}

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}
 

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()))

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))
 

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)

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) {