wip

.gitignore

@@ -39,3 +39,4 @@ prime/
 .snapcraft/
 yq*.snap
 test.yml
+test2.yml

acceptance_tests/leading-seperator.sh

@@ -7,6 +7,25 @@ a: test
 EOL
 }
 
+testLeadingSeperatorWithDoc() {
+  cat >test.yml <<EOL
+# hi peeps
+# cool
+---
+a: test
+EOL
+
+  read -r -d '' expected << EOM
+# hi peeps
+# cool
+---
+a: thing
+EOM
+
+  X=$(./yq e '.a = "thing"' - < test.yml)
+  assertEquals "$expected" "$X"
+}
+
 testLeadingSeperatorPipeIntoEvalSeq() {
   X=$(./yq e - < test.yml)
   expected=$(cat test.yml)
@@ -33,7 +52,7 @@ testLeadingSeperatorEvalAll() {
   assertEquals "$expected" "$X"
 }
 
-testLeadingSeperatorMultiDocEval() {
+testLeadingSeperatorMultiDocEvalSimple() {
   read -r -d '' expected << EOM
 ---
 a: test
@@ -47,6 +66,111 @@ EOM
   assertEquals "$expected" "$X"
 }
 
+testLeadingSeperatorMultiDocInOneFile() {
+  cat >test.yml <<EOL
+---
+# hi peeps
+# cool
+a: test
+---
+b: things
+EOL
+  expected=$(cat test.yml)
+  X=$(./yq e '.' test.yml)
+  assertEquals "$expected" "$X"
+}
+
+testLeadingSeperatorMultiDocEvalComments() {
+  cat >test.yml <<EOL
+# hi peeps
+# cool
+a: test
+EOL
+
+cat >test2.yml <<EOL
+# this is another doc
+# great
+b: sane
+EOL
+
+  read -r -d '' expected << EOM
+# hi peeps
+# cool
+a: test
+---
+# this is another doc
+# great
+b: sane
+EOM
+
+
+  X=$(./yq e '.' test.yml test2.yml)
+  assertEquals "$expected" "$X"
+}
+
+testLeadingSeperatorMultiDocEvalCommentsTrailingSep() {
+  cat >test.yml <<EOL
+# hi peeps
+# cool
+---
+a: test
+EOL
+
+cat >test2.yml <<EOL
+# this is another doc
+# great
+---
+b: sane
+EOL
+
+  read -r -d '' expected << EOM
+# hi peeps
+# cool
+---
+a: test
+---
+# this is another doc
+# great
+---
+b: sane
+EOM
+
+
+  X=$(./yq e '.' test.yml test2.yml)
+  assertEquals "$expected" "$X"
+}
+
+testLeadingSeperatorMultiDocEvalCommentsLeadingSep() {
+  cat >test.yml <<EOL
+---
+# hi peeps
+# cool
+a: test
+EOL
+
+cat >test2.yml <<EOL
+---
+# this is another doc
+# great
+b: sane
+EOL
+
+  read -r -d '' expected << EOM
+---
+# hi peeps
+# cool
+a: test
+---
+# this is another doc
+# great
+b: sane
+EOM
+
+
+  X=$(./yq e '.' test.yml test2.yml)
+  assertEquals "$expected" "$X"
+}
+
 testLeadingSeperatorMultiDocEvalAll() {
   read -r -d '' expected << EOM
 ---

cmd/evaluate_all_command.go

@@ -119,7 +119,7 @@ func evaluateAll(cmd *cobra.Command, args []string) error {
 		}
 	case 1:
 		if nullInput {
-			err = yqlib.NewStreamEvaluator().EvaluateNew(processExpression(args[0]), printer)
+			err = yqlib.NewStreamEvaluator().EvaluateNew(processExpression(args[0]), printer, "")
 		} else {
 			err = allAtOnceEvaluator.EvaluateFiles(processExpression(""), []string{args[0]}, printer)
 		}

cmd/evalute_sequence_command.go

@@ -132,7 +132,7 @@ func evaluateSequence(cmd *cobra.Command, args []string) error {
 		}
 	case 1:
 		if nullInput {
-			err = streamEvaluator.EvaluateNew(processExpression(args[0]), printer)
+			err = streamEvaluator.EvaluateNew(processExpression(args[0]), printer, "")
 		} else {
 			err = streamEvaluator.EvaluateFiles(processExpression(""), []string{args[0]}, printer)
 		}

pkg/yqlib/all_at_once_evaluator.go

@@ -2,7 +2,6 @@ package yqlib
 
 import (
 	"container/list"
-	"os"
 
 	yaml "gopkg.in/yaml.v3"
 )
@@ -49,17 +48,17 @@ func (e *allAtOnceEvaluator) EvaluateCandidateNodes(expression string, inputCand
 
 func (e *allAtOnceEvaluator) EvaluateFiles(expression string, filenames []string, printer Printer) error {
 	fileIndex := 0
-	firstFileLeadingSeperator := false
+	firstFileLeadingContent := ""
 
 	var allDocuments *list.List = list.New()
 	for _, filename := range filenames {
-		reader, leadingSeperator, err := readStream(filename)
+		reader, leadingContent, err := readStream(filename)
 		if err != nil {
 			return err
 		}
 
-		if fileIndex == 0 && leadingSeperator {
+		if fileIndex == 0 {
-			firstFileLeadingSeperator = leadingSeperator
+			firstFileLeadingContent = leadingContent
 		}
 
 		fileDocuments, err := readDocuments(reader, filename, fileIndex)
@@ -78,24 +77,11 @@ func (e *allAtOnceEvaluator) EvaluateFiles(expression string, filenames []string
 			FileIndex: 0,
 		}
 		allDocuments.PushBack(candidateNode)
-
-		if len(filenames) > 0 {
-			reader, _, err := readStream(filenames[0])
-			if err != nil {
-				return err
-			}
-			switch reader := reader.(type) {
-			case *os.File:
-				defer safelyCloseFile(reader)
-			}
-			printer.SetPreamble(reader)
-		}
 	}
 
 	matches, err := e.EvaluateCandidateNodes(expression, allDocuments)
 	if err != nil {
 		return err
 	}
-	printer.SetPrintLeadingSeperator(firstFileLeadingSeperator)
+	return printer.PrintResults(matches, firstFileLeadingContent)
-	return printer.PrintResults(matches)
 }

pkg/yqlib/operators_test.go

@@ -256,7 +256,7 @@ func documentOutput(t *testing.T, w *bufio.Writer, s expressionScenario, formatt
 		t.Error(err, s.expression)
 	}
 
-	err = printer.PrintResults(context.MatchingNodes)
+	err = printer.PrintResults(context.MatchingNodes, "")
 	if err != nil {
 		t.Error(err, s.expression)
 	}

pkg/yqlib/printer.go

@@ -4,17 +4,14 @@ import (
 	"bufio"
 	"container/list"
 	"io"
+	"strings"
 
 	yaml "gopkg.in/yaml.v3"
 )
 
 type Printer interface {
-	PrintResults(matchingNodes *list.List) error
+	PrintResults(matchingNodes *list.List, leadingContent string) error
 	PrintedAnything() bool
-	SetPrintLeadingSeperator(bool)
-
-	SetPreamble(reader io.Reader)
-
 	//e.g. when given a front-matter doc, like jekyll
 	SetAppendix(reader io.Reader)
 }
@@ -48,13 +45,6 @@ func NewPrinter(writer io.Writer, outputToJSON bool, unwrapScalar bool, colorsEn
 	}
 }
 
-func (p *resultsPrinter) SetPrintLeadingSeperator(printLeadingSeperator bool) {
-	if printLeadingSeperator {
-		p.firstTimePrinting = false
-		p.previousFileIndex = -1
-	}
-}
-
 func (p *resultsPrinter) SetPreamble(reader io.Reader) {
 	p.preambleReader = reader
 }
@@ -95,7 +85,7 @@ func (p *resultsPrinter) safelyFlush(writer *bufio.Writer) {
 	}
 }
 
-func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
+func (p *resultsPrinter) PrintResults(matchingNodes *list.List, leadingContent string) error {
 	log.Debug("PrintResults for %v matches", matchingNodes.Len())
 	if p.outputToJSON {
 		explodeOp := Operation{OperationType: explodeOpType}
@@ -110,15 +100,8 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 	bufferedWriter := bufio.NewWriter(p.writer)
 	defer p.safelyFlush(bufferedWriter)
 
-	if p.preambleReader != nil && !p.outputToJSON {
-		log.Debug("Piping preamble reader...")
-		_, err := io.Copy(bufferedWriter, p.preambleReader)
-		if err != nil {
-			return err
-		}
-	}
-
 	if matchingNodes.Len() == 0 {
+		p.writeString(bufferedWriter, leadingContent)
 		log.Debug("no matching results, nothing to print")
 		return nil
 	}
@@ -129,16 +112,25 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 		p.firstTimePrinting = false
 	}
 
+	printedLead := false
+
 	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)
-		if (p.previousDocIndex != mappedDoc.Document || p.previousFileIndex != mappedDoc.FileIndex) && p.printDocSeparators {
+		if (p.previousDocIndex != mappedDoc.Document || p.previousFileIndex != mappedDoc.FileIndex) && p.printDocSeparators &&
+			(printedLead || !strings.HasPrefix(leadingContent, "---")) {
 			log.Debug("-- writing doc sep")
 			if err := p.writeString(bufferedWriter, "---\n"); err != nil {
 				return err
 			}
 		}
 
+		if !printedLead {
+			// we want to print this after the seperator logic
+			p.writeString(bufferedWriter, leadingContent)
+			printedLead = true
+		}
+
 		if err := p.printNode(mappedDoc.Node, bufferedWriter); err != nil {
 			return err
 		}

pkg/yqlib/stream_evaluator.go

@@ -12,9 +12,9 @@ import (
 // Uses less memory than loading all documents and running the expression once, but this cannot process
 // cross document expressions.
 type StreamEvaluator interface {
-	Evaluate(filename string, reader io.Reader, node *ExpressionNode, printer Printer) (uint, error)
+	Evaluate(filename string, reader io.Reader, node *ExpressionNode, printer Printer, leadingContent string) (uint, error)
 	EvaluateFiles(expression string, filenames []string, printer Printer) error
-	EvaluateNew(expression string, printer Printer) error
+	EvaluateNew(expression string, printer Printer, leadingContent string) error
 }
 
 type streamEvaluator struct {
@@ -27,7 +27,7 @@ func NewStreamEvaluator() StreamEvaluator {
 	return &streamEvaluator{treeNavigator: NewDataTreeNavigator(), treeCreator: NewExpressionParser()}
 }
 
-func (s *streamEvaluator) EvaluateNew(expression string, printer Printer) error {
+func (s *streamEvaluator) EvaluateNew(expression string, printer Printer, leadingContent string) error {
 	node, err := s.treeCreator.ParseExpression(expression)
 	if err != nil {
 		return err
@@ -45,7 +45,7 @@ func (s *streamEvaluator) EvaluateNew(expression string, printer Printer) error
 	if errorParsing != nil {
 		return errorParsing
 	}
-	return printer.PrintResults(result.MatchingNodes)
+	return printer.PrintResults(result.MatchingNodes, leadingContent)
 }
 
 func (s *streamEvaluator) EvaluateFiles(expression string, filenames []string, printer Printer) error {
@@ -55,15 +55,19 @@ func (s *streamEvaluator) EvaluateFiles(expression string, filenames []string, p
 		return err
 	}
 
+	var firstFileLeadingContent string
+
 	for index, filename := range filenames {
-		reader, leadingSeperator, err := readStream(filename)
+		reader, leadingContent, err := readStream(filename)
-		if index == 0 && leadingSeperator {
+
-			printer.SetPrintLeadingSeperator(leadingSeperator)
+		if index == 0 {
+			firstFileLeadingContent = leadingContent
 		}
+
 		if err != nil {
 			return err
 		}
-		processedDocs, err := s.Evaluate(filename, reader, node, printer)
+		processedDocs, err := s.Evaluate(filename, reader, node, printer, leadingContent)
 		if err != nil {
 			return err
 		}
@@ -76,26 +80,13 @@ func (s *streamEvaluator) EvaluateFiles(expression string, filenames []string, p
 	}
 
 	if totalProcessDocs == 0 {
-
+		return s.EvaluateNew(expression, printer, firstFileLeadingContent)
-		if len(filenames) > 0 {
-			reader, _, err := readStream(filenames[0])
-			if err != nil {
-				return err
-			}
-			switch reader := reader.(type) {
-			case *os.File:
-				defer safelyCloseFile(reader)
-			}
-			printer.SetPreamble(reader)
-		}
-
-		return s.EvaluateNew(expression, printer)
 	}
 
 	return nil
 }
 
-func (s *streamEvaluator) Evaluate(filename string, reader io.Reader, node *ExpressionNode, printer Printer) (uint, error) {
+func (s *streamEvaluator) Evaluate(filename string, reader io.Reader, node *ExpressionNode, printer Printer, leadingContent string) (uint, error) {
 
 	var currentIndex uint
 	decoder := yaml.NewDecoder(reader)
@@ -122,7 +113,13 @@ func (s *streamEvaluator) Evaluate(filename string, reader io.Reader, node *Expr
 		if errorParsing != nil {
 			return currentIndex, errorParsing
 		}
-		err := printer.PrintResults(result.MatchingNodes)
+		var err error
+		if currentIndex == 0 {
+			err = printer.PrintResults(result.MatchingNodes, leadingContent)
+		} else {
+			err = printer.PrintResults(result.MatchingNodes, "")
+		}
+
 		if err != nil {
 			return currentIndex, err
 		}

pkg/yqlib/utils.go

@@ -5,42 +5,49 @@ import (
 	"container/list"
 	"io"
 	"os"
+	"regexp"
+	"strings"
 
 	yaml "gopkg.in/yaml.v3"
 )
 
-func readStream(filename string) (io.Reader, bool, error) {
+func readStream(filename string) (io.Reader, string, error) {
-
+	var commentLineRegEx = regexp.MustCompile(`^\s*#`)
+	var reader *bufio.Reader
 	if filename == "-" {
-		reader := bufio.NewReader(os.Stdin)
+		reader = bufio.NewReader(os.Stdin)
-
-		seperatorBytes, err := reader.Peek(3)
-
-		if err == io.EOF {
-			// EOF are handled else where..
-			return reader, false, nil
-		}
-
-		return reader, string(seperatorBytes) == "---", err
 	} else {
 		// ignore CWE-22 gosec issue - that's more targetted for http based apps that run in a public directory,
 		// and ensuring that it's not possible to give a path to a file outside thar directory.
-		reader, err := os.Open(filename) // #nosec
+		file, err := os.Open(filename) // #nosec
 		if err != nil {
-			return nil, false, err
+			return nil, "", err
 		}
-		seperatorBytes := make([]byte, 3)
+		reader = bufio.NewReader(file)
-		_, err = reader.Read(seperatorBytes)
+	}
+	var sb strings.Builder
+
+	for {
+		peekBytes, err := reader.Peek(3)
+
 		if err == io.EOF {
 			// EOF are handled else where..
-			return reader, false, nil
+			return reader, sb.String(), nil
 		} else if err != nil {
-			return nil, false, err
+			return reader, sb.String(), err
+		} else if string(peekBytes) == "---" || commentLineRegEx.MatchString(string(peekBytes)) {
+			line, err := reader.ReadString('\n')
+			sb.WriteString(line)
+			if err == io.EOF {
+				return reader, sb.String(), nil
+			} else if err != nil {
+				return reader, sb.String(), err
+			}
+		} else {
+			return reader, sb.String(), nil
+		}
 	}
-		_, err = reader.Seek(0, 0)
 
-		return reader, string(seperatorBytes) == "---", err
-	}
 }
 
 func readDocuments(reader io.Reader, filename string, fileIndex int) (*list.List, error) {