Split printer!

pkg/yqlib/operators_test.go

@@ -107,7 +107,7 @@ func copyFromHeader(title string, out *os.File) error {
 
 func formatYaml(yaml string, filename string) string {
 	var output bytes.Buffer
-	printer := NewPrinter(bufio.NewWriter(&output), YamlOutputFormat, true, false, 2, true)
+	printer := NewPrinterWithSingleWriter(bufio.NewWriter(&output), YamlOutputFormat, true, false, 2, true)
 
 	node, err := NewExpressionParser().ParseExpression(".. style= \"\"")
 	if err != nil {
@@ -216,7 +216,7 @@ func documentInput(w *bufio.Writer, s expressionScenario) (string, string) {
 func documentOutput(t *testing.T, w *bufio.Writer, s expressionScenario, formattedDoc string, formattedDoc2 string) {
 	var output bytes.Buffer
 	var err error
-	printer := NewPrinter(bufio.NewWriter(&output), YamlOutputFormat, true, false, 2, true)
+	printer := NewPrinterWithSingleWriter(bufio.NewWriter(&output), YamlOutputFormat, true, false, 2, true)
 
 	node, err := NewExpressionParser().ParseExpression(s.expression)
 	if err != nil {

pkg/yqlib/printer.go

@@ -102,13 +102,6 @@ func (p *resultsPrinter) writeString(writer io.Writer, txt string) error {
 	return errorWriting
 }
 
-func (p *resultsPrinter) safelyFlush(writer *bufio.Writer) {
-	if err := writer.Flush(); err != nil {
-		log.Error("Error flushing writer!")
-		log.Error(err.Error())
-	}
-}
-
 func (p *resultsPrinter) processLeadingContent(mappedDoc *CandidateNode, writer io.Writer) error {
 	if strings.Contains(mappedDoc.Node.HeadComment, "$yqLeadingContent$") {
 		log.Debug("headcommentwas %v", mappedDoc.Node.HeadComment)
@@ -177,14 +170,12 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 		p.firstTimePrinting = false
 	}
 
-	index := 0
-
 	for el := matchingNodes.Front(); el != nil; el = el.Next() {
 
 		mappedDoc := el.Value.(*CandidateNode)
 		log.Debug("-- print sep logic: p.firstTimePrinting: %v, previousDocIndex: %v, mappedDoc.Document: %v, printDocSeparators: %v", p.firstTimePrinting, p.previousDocIndex, mappedDoc.Document, p.printDocSeparators)
 
-		writer, errorWriting := p.printerWriter.GetWriter(mappedDoc, index)
+		writer, errorWriting := p.printerWriter.GetWriter(mappedDoc)
 		if errorWriting != nil {
 			return errorWriting
 		}
@@ -210,13 +201,10 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 		if err := writer.Flush(); err != nil {
 			return err
 		}
-
-		index++
-
 	}
 
 	if p.appendixReader != nil && p.outputFormat == YamlOutputFormat {
-		writer, err := p.printerWriter.GetWriter(nil, index)
+		writer, err := p.printerWriter.GetWriter(nil)
 		if err != nil {
 			return err
 		}
@@ -227,6 +215,9 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 		if err != nil {
 			return err
 		}
+		if err := writer.Flush(); err != nil {
+			return err
+		}
 	}
 
 	return nil

pkg/yqlib/printer_writer.go

@@ -5,10 +5,12 @@ import (
 	"fmt"
 	"io"
 	"os"
+
+	"gopkg.in/yaml.v3"
 )
 
 type PrinterWriter interface {
-	GetWriter(node *CandidateNode, index int) (*bufio.Writer, error)
+	GetWriter(node *CandidateNode) (*bufio.Writer, error)
 }
 
 type singlePrinterWriter struct {
@@ -21,7 +23,7 @@ func NewSinglePrinterWriter(writer io.Writer) PrinterWriter {
 	}
 }
 
-func (sp *singlePrinterWriter) GetWriter(node *CandidateNode, i int) (*bufio.Writer, error) {
+func (sp *singlePrinterWriter) GetWriter(node *CandidateNode) (*bufio.Writer, error) {
 	return sp.bufferedWriter, nil
 }
 
@@ -29,6 +31,7 @@ type multiPrintWriter struct {
 	treeNavigator  DataTreeNavigator
 	nameExpression *ExpressionNode
 	extension      string
+	index          int
 }
 
 func NewMultiPrinterWriter(expression *ExpressionNode, format PrinterOutputFormat) PrinterWriter {
@@ -45,33 +48,33 @@ func NewMultiPrinterWriter(expression *ExpressionNode, format PrinterOutputForma
 		nameExpression: expression,
 		extension:      extension,
 		treeNavigator:  NewDataTreeNavigator(),
+		index:          0,
 	}
 }
 
-func (sp *multiPrintWriter) GetWriter(node *CandidateNode, index int) (*bufio.Writer, error) {
+func (sp *multiPrintWriter) GetWriter(node *CandidateNode) (*bufio.Writer, error) {
 	name := ""
 
-	if sp.nameExpression != nil {
-		context := Context{MatchingNodes: node.AsList()}
-		result, err := sp.treeNavigator.GetMatchingNodes(context, sp.nameExpression)
-		if err != nil {
-			return nil, err
-		}
-		if result.MatchingNodes.Len() > 0 {
-			name = result.MatchingNodes.Front().Value.(*CandidateNode).Node.Value
-		}
+	indexVariableNode := yaml.Node{Kind: yaml.ScalarNode, Tag: "!!int", Value: fmt.Sprintf("%v", sp.index)}
+	indexVariableCandidate := CandidateNode{Node: &indexVariableNode}
+
+	context := Context{MatchingNodes: node.AsList()}
+	context.SetVariable("index", indexVariableCandidate.AsList())
+	result, err := sp.treeNavigator.GetMatchingNodes(context, sp.nameExpression)
+	if err != nil {
+		return nil, err
 	}
-	if name == "" {
-		name = fmt.Sprintf("%v.%v", index, sp.extension)
-	} else {
-		name = fmt.Sprintf("%v.%v", name, sp.extension)
+	if result.MatchingNodes.Len() > 0 {
+		name = result.MatchingNodes.Front().Value.(*CandidateNode).Node.Value
 	}
+	name = fmt.Sprintf("%v.%v", name, sp.extension)
 
 	f, err := os.Create(name)
 
 	if err != nil {
 		return nil, err
 	}
+	sp.index = sp.index + 1
 
 	return bufio.NewWriter(f), nil