From 54b20f68dbb04f6d084ef905e05ead76b8642980 Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Fri, 29 Oct 2021 16:07:17 +1100
Subject: [PATCH] Wip split printer

---
 pkg/yqlib/printer.go        | 40 ++++++++++++++++++----------
 pkg/yqlib/printer_writer.go | 52 +++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 13 deletions(-)
 create mode 100644 pkg/yqlib/printer_writer.go

diff --git a/pkg/yqlib/printer.go b/pkg/yqlib/printer.go
index b3b07dc5..6685e68a 100644
--- a/pkg/yqlib/printer.go
+++ b/pkg/yqlib/printer.go
@@ -44,7 +44,7 @@ type resultsPrinter struct {
 	colorsEnabled      bool
 	indent             int
 	printDocSeparators bool
-	writer             io.Writer
+	printerWriter      printerWriter
 	firstTimePrinting  bool
 	previousDocIndex   uint
 	previousFileIndex  int
@@ -53,9 +53,9 @@ type resultsPrinter struct {
 	appendixReader     io.Reader
 }
 
-func NewPrinter(writer io.Writer, outputFormat PrinterOutputFormat, unwrapScalar bool, colorsEnabled bool, indent int, printDocSeparators bool) Printer {
+func NewPrinter(printerWriter printerWriter, outputFormat PrinterOutputFormat, unwrapScalar bool, colorsEnabled bool, indent int, printDocSeparators bool) Printer {
 	return &resultsPrinter{
-		writer:             writer,
+		printerWriter:      printerWriter,
 		outputFormat:       outputFormat,
 		unwrapScalar:       unwrapScalar,
 		colorsEnabled:      colorsEnabled,
@@ -150,6 +150,12 @@ func (p *resultsPrinter) processLeadingContent(mappedDoc *CandidateNode, writer
 
 func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 	log.Debug("PrintResults for %v matches", matchingNodes.Len())
+
+	if matchingNodes.Len() == 0 {
+		log.Debug("no matching results, nothing to print")
+		return nil
+	}
+
 	if p.outputFormat != YamlOutputFormat {
 		explodeOp := Operation{OperationType: explodeOpType}
 		explodeNode := ExpressionNode{Operation: &explodeOp}
@@ -160,10 +166,6 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 		matchingNodes = context.MatchingNodes
 	}
 
-	if matchingNodes.Len() == 0 {
-		log.Debug("no matching results, nothing to print")
-		return nil
-	}
 	if p.firstTimePrinting {
 		node := matchingNodes.Front().Value.(*CandidateNode)
 		p.previousDocIndex = node.Document
@@ -171,37 +173,49 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 		p.firstTimePrinting = false
 	}
 
-	bufferedWriter := bufio.NewWriter(p.writer)
-	defer p.safelyFlush(bufferedWriter)
+	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)
+		if errorWriting != nil {
+			return errorWriting
+		}
+
 		commentStartsWithSeparator := strings.Contains(mappedDoc.Node.HeadComment, "$yqLeadingContent$\n$yqDocSeperator$")
 
 		if (p.previousDocIndex != mappedDoc.Document || p.previousFileIndex != mappedDoc.FileIndex) && p.printDocSeparators && !commentStartsWithSeparator {
 			log.Debug("-- writing doc sep")
-			if err := p.writeString(bufferedWriter, "---\n"); err != nil {
+			if err := p.writeString(writer, "---\n"); err != nil {
 				return err
 			}
 		}
 
-		if err := p.processLeadingContent(mappedDoc, bufferedWriter); err != nil {
+		if err := p.processLeadingContent(mappedDoc, writer); err != nil {
 			return err
 		}
 
-		if err := p.printNode(mappedDoc.Node, bufferedWriter); err != nil {
+		if err := p.printNode(mappedDoc.Node, writer); err != nil {
 			return err
 		}
 
 		p.previousDocIndex = mappedDoc.Document
+		if err := writer.Flush(); err != nil {
+			return err
+		}
+
+		index++
+
 	}
 
 	if p.appendixReader != nil && p.outputFormat == YamlOutputFormat {
+		writer := p.printerWriter.GetWriter(nil, index)
 		log.Debug("Piping appendix reader...")
 		betterReader := bufio.NewReader(p.appendixReader)
-		_, err := io.Copy(bufferedWriter, betterReader)
+		_, err := io.Copy(writer, betterReader)
 		if err != nil {
 			return err
 		}
diff --git a/pkg/yqlib/printer_writer.go b/pkg/yqlib/printer_writer.go
new file mode 100644
index 00000000..30d05f5b
--- /dev/null
+++ b/pkg/yqlib/printer_writer.go
@@ -0,0 +1,52 @@
+package yqlib
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+)
+
+type printerWriter interface {
+	GetWriter(node *CandidateNode, index int) (*bufio.Writer, error)
+}
+
+type singlePrinterWriter struct {
+	bufferedWriter *bufio.Writer
+}
+
+func NewSinglePrinterWriter(writer io.Writer) printerWriter {
+	return &singlePrinterWriter{
+		bufferedWriter: bufio.NewWriter(writer),
+	}
+}
+
+func (sp *singlePrinterWriter) GetWriter(node *CandidateNode, i int) (*bufio.Writer, error) {
+	return sp.bufferedWriter, nil
+}
+
+type multiPrintWriter struct {
+	treeNavigator  DataTreeNavigator
+	nameExpression *ExpressionNode
+	extension      string
+}
+
+func (sp *multiPrintWriter) GetWriter(node *CandidateNode, index int) (*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
+		}
+	}
+	if name == "" {
+		name = fmt.Sprintf("%v.%v", index, sp.extension)
+	} else {
+		name = fmt.Sprintf("%v.%v", name, sp.extension)
+	}
+
+}