From 8ad8903e21a9a42dcbc1aa2a00e6e797d00e531d Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Sat, 30 Oct 2021 09:56:47 +1100
Subject: [PATCH] Split printer!

---
 .gitignore                            |   6 +-
 acceptance_tests/basic.sh             |   2 +-
 acceptance_tests/empty.sh             |   1 +
 acceptance_tests/front-matter.sh      |   1 +
 acceptance_tests/leading-seperator.sh |   1 +
 acceptance_tests/output-format.sh     |   4 +
 acceptance_tests/pretty-print.sh      |   3 +
 acceptance_tests/split-printer.sh     | 157 ++++++++++++++++++++++++++
 cmd/root.go                           |   2 +-
 go.mod                                |   5 +
 pkg/yqlib/operators_test.go           |   4 +-
 pkg/yqlib/printer.go                  |  19 +---
 pkg/yqlib/printer_writer.go           |  35 +++---
 13 files changed, 205 insertions(+), 35 deletions(-)
 create mode 100755 acceptance_tests/split-printer.sh

diff --git a/.gitignore b/.gitignore
index fac7e030..00b1ab74 100644
--- a/.gitignore
+++ b/.gitignore
@@ -38,8 +38,12 @@ parts/
 prime/
 .snapcraft/
 yq*.snap
+
 test.yml
-test2.yml
+test*.yml
+0.yml
+1.yml
+2.yml
 
 
 # man page
diff --git a/acceptance_tests/basic.sh b/acceptance_tests/basic.sh
index f83f9085..dc9069ae 100755
--- a/acceptance_tests/basic.sh
+++ b/acceptance_tests/basic.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 
 setUp() {
-  rm -f test.yml
+  rm test*.yml || true
 }
 
 testBasicEvalRoundTrip() {
diff --git a/acceptance_tests/empty.sh b/acceptance_tests/empty.sh
index d7f344df..7ffb813f 100755
--- a/acceptance_tests/empty.sh
+++ b/acceptance_tests/empty.sh
@@ -1,6 +1,7 @@
 #!/bin/bash
 
 setUp() {
+  rm test*.yml || true
   cat >test.yml <<EOL
 # comment
 EOL
diff --git a/acceptance_tests/front-matter.sh b/acceptance_tests/front-matter.sh
index b02356b5..a9f079be 100755
--- a/acceptance_tests/front-matter.sh
+++ b/acceptance_tests/front-matter.sh
@@ -1,6 +1,7 @@
 #!/bin/bash
 
 setUp() {
+  rm test*.yml || true
   cat >test.yml <<EOL
 ---
 a: apple
diff --git a/acceptance_tests/leading-seperator.sh b/acceptance_tests/leading-seperator.sh
index a8fae4e9..65338b36 100755
--- a/acceptance_tests/leading-seperator.sh
+++ b/acceptance_tests/leading-seperator.sh
@@ -1,6 +1,7 @@
 #!/bin/bash
 
 setUp() {
+  rm test*.yml || true
   cat >test.yml <<EOL
 ---
 a: test
diff --git a/acceptance_tests/output-format.sh b/acceptance_tests/output-format.sh
index a36290b5..35e3b92f 100755
--- a/acceptance_tests/output-format.sh
+++ b/acceptance_tests/output-format.sh
@@ -1,5 +1,9 @@
 #!/bin/bash
 
+setUp() {
+  rm test*.yml || true
+}
+
 testOutputJsonDeprecated() {
   cat >test.yml <<EOL
 a: {b: ["cat"]}
diff --git a/acceptance_tests/pretty-print.sh b/acceptance_tests/pretty-print.sh
index 5ebfed21..9b59ea04 100755
--- a/acceptance_tests/pretty-print.sh
+++ b/acceptance_tests/pretty-print.sh
@@ -1,5 +1,8 @@
 #!/bin/bash
 
+setUp() {
+  rm test*.yml || true
+}
 
 testPrettyPrintWithBooleans() {
   cat >test.yml <<EOL
diff --git a/acceptance_tests/split-printer.sh b/acceptance_tests/split-printer.sh
new file mode 100755
index 00000000..71170422
--- /dev/null
+++ b/acceptance_tests/split-printer.sh
@@ -0,0 +1,157 @@
+#!/bin/bash
+
+setUp() {
+  rm test*.yml || true
+}
+
+testBasicSplitWithName() {
+  cat >test.yml <<EOL
+a: test_doc1
+--- 
+a: test_doc2
+EOL
+
+  ./yq e test.yml -s ".a"
+
+  doc1=$(cat test_doc1.yml)
+  
+  assertEquals "a: test_doc1" "$doc1"
+
+  doc2=$(cat test_doc2.yml)
+  read -r -d '' expectedDoc2 << EOM
+---
+a: test_doc2
+EOM
+  assertEquals "$expectedDoc2" "$doc2"
+}
+
+testBasicSplitWithNameEvalAll() {
+  cat >test.yml <<EOL
+a: test_doc1
+--- 
+a: test_doc2
+EOL
+
+  ./yq ea test.yml -s ".a"
+
+  doc1=$(cat test_doc1.yml)
+  
+  assertEquals "a: test_doc1" "$doc1"
+
+  doc2=$(cat test_doc2.yml)
+  read -r -d '' expectedDoc2 << EOM
+---
+a: test_doc2
+EOM
+  assertEquals "$expectedDoc2" "$doc2"
+}
+
+testBasicSplitWithIndex() {
+  cat >test.yml <<EOL
+a: test_doc1
+--- 
+a: test_doc2
+EOL
+
+  ./yq e test.yml -s '"test_" + $index'
+
+  doc1=$(cat test_0.yml)
+  
+  assertEquals "a: test_doc1" "$doc1"
+
+  doc2=$(cat test_1.yml)
+  read -r -d '' expectedDoc2 << EOM
+---
+a: test_doc2
+EOM
+  assertEquals "$expectedDoc2" "$doc2"
+}
+
+testBasicSplitWithIndexEvalAll() {
+  cat >test.yml <<EOL
+a: test_doc1
+--- 
+a: test_doc2
+EOL
+
+  ./yq ea test.yml -s '"test_" + $index'
+
+  doc1=$(cat test_0.yml)
+  
+  assertEquals "a: test_doc1" "$doc1"
+
+  doc2=$(cat test_1.yml)
+  read -r -d '' expectedDoc2 << EOM
+---
+a: test_doc2
+EOM
+  assertEquals "$expectedDoc2" "$doc2"
+}
+
+
+testArraySplitWithNameNoSeparators() {
+  cat >test.yml <<EOL
+- name: test_fred
+  age: 35
+- name: test_catherine
+  age: 37
+EOL
+
+  ./yq e --no-doc -s ".name"  ".[]" test.yml 
+
+  doc1=$(cat test_fred.yml)
+  read -r -d '' expectedDoc1 << EOM
+name: test_fred
+age: 35
+EOM
+
+  assertEquals "$expectedDoc1" "$doc1"
+
+  doc2=$(cat test_catherine.yml)
+  read -r -d '' expectedDoc2 << EOM
+name: test_catherine
+age: 37
+EOM
+  assertEquals "$expectedDoc2" "$doc2"
+}
+
+testArraySplitWithNameNoSeparatorsEvalAll() {
+  cat >test.yml <<EOL
+- name: test_fred
+  age: 35
+- name: test_catherine
+  age: 37
+EOL
+
+cat >test2.yml <<EOL
+- name: test_mike
+  age: 564
+EOL
+
+  ./yq ea --no-doc -s ".name"  ".[]" test.yml test2.yml
+
+  doc1=$(cat test_fred.yml)
+  read -r -d '' expectedDoc1 << EOM
+name: test_fred
+age: 35
+EOM
+
+  assertEquals "$expectedDoc1" "$doc1"
+
+  doc2=$(cat test_catherine.yml)
+  read -r -d '' expectedDoc2 << EOM
+name: test_catherine
+age: 37
+EOM
+  assertEquals "$expectedDoc2" "$doc2"
+
+
+  doc3=$(cat test_mike.yml)
+  read -r -d '' expectedDoc3 << EOM
+name: test_mike
+age: 564
+EOM
+  assertEquals "$expectedDoc3" "$doc3"
+}
+
+source ./scripts/shunit2
\ No newline at end of file
diff --git a/cmd/root.go b/cmd/root.go
index 1a8fb613..11471d27 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -64,7 +64,7 @@ See https://mikefarah.gitbook.io/yq/ for detailed documentation and examples.`,
 	rootCmd.PersistentFlags().StringVarP(&frontMatter, "front-matter", "f", "", "(extract|process) first input as yaml front-matter. Extract will pull out the yaml content, process will run the expression against the yaml content, leaving the remaining data intact")
 	rootCmd.PersistentFlags().BoolVarP(&leadingContentPreProcessing, "header-preprocess", "", true, "Slurp any header comments and separators before processing expression. This is a workaround for go-yaml to persist header content properly.")
 
-	rootCmd.PersistentFlags().StringVarP(&splitFileExp, "split-exp", "s", "", "print each result (or doc) into a file named (exp). Set to 'i' if you'd like to have numbered output.")
+	rootCmd.PersistentFlags().StringVarP(&splitFileExp, "split-exp", "s", "", "print each result (or doc) into a file named (exp). [exp] argument must return a string. You can use $index in the expression as the result counter.")
 
 	rootCmd.AddCommand(
 		createEvaluateSequenceCommand(),
diff --git a/go.mod b/go.mod
index bc4fa628..6f7e79f4 100644
--- a/go.mod
+++ b/go.mod
@@ -13,8 +13,13 @@ require (
 )
 
 require (
+	github.com/inconshreveable/mousetrap v1.0.0 // indirect
+	github.com/mattn/go-colorable v0.1.8 // indirect
+	github.com/mattn/go-isatty v0.0.12 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/timtadh/data-structures v0.5.3 // indirect
 	golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e // indirect
+	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
 )
 
 go 1.17
diff --git a/pkg/yqlib/operators_test.go b/pkg/yqlib/operators_test.go
index b450e0c6..af7b1d7a 100644
--- a/pkg/yqlib/operators_test.go
+++ b/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 {
diff --git a/pkg/yqlib/printer.go b/pkg/yqlib/printer.go
index ac95e943..3ac4bcee 100644
--- a/pkg/yqlib/printer.go
+++ b/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
diff --git a/pkg/yqlib/printer_writer.go b/pkg/yqlib/printer_writer.go
index d74baf52..d4292323 100644
--- a/pkg/yqlib/printer_writer.go
+++ b/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