Split printer!

.gitignore

@@ -38,8 +38,12 @@ parts/
 prime/
 .snapcraft/
 yq*.snap
+
 test.yml
-test2.yml
+test*.yml
+0.yml
+1.yml
+2.yml
 
 
 # man page

acceptance_tests/basic.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
 setUp() {
-  rm -f test.yml
+  rm test*.yml || true
 }
 
 testBasicEvalRoundTrip() {

acceptance_tests/empty.sh

@@ -1,6 +1,7 @@
 #!/bin/bash
 
 setUp() {
+  rm test*.yml || true
   cat >test.yml <<EOL
 # comment
 EOL

acceptance_tests/front-matter.sh

@@ -1,6 +1,7 @@
 #!/bin/bash
 
 setUp() {
+  rm test*.yml || true
   cat >test.yml <<EOL
 ---
 a: apple

acceptance_tests/leading-seperator.sh

@@ -1,6 +1,7 @@
 #!/bin/bash
 
 setUp() {
+  rm test*.yml || true
   cat >test.yml <<EOL
 ---
 a: test

acceptance_tests/output-format.sh

@@ -1,5 +1,9 @@
 #!/bin/bash
 
+setUp() {
+  rm test*.yml || true
+}
+
 testOutputJsonDeprecated() {
   cat >test.yml <<EOL
 a: {b: ["cat"]}

acceptance_tests/pretty-print.sh

@@ -1,5 +1,8 @@
 #!/bin/bash
 
+setUp() {
+  rm test*.yml || true
+}
 
 testPrettyPrintWithBooleans() {
   cat >test.yml <<EOL

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

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

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

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