Added front-matter handler

pkg/yqlib/file_utils.go

@@ -2,6 +2,7 @@ package yqlib
 
 import (
 	"io"
+	"io/ioutil"
 	"os"
 )
 
@@ -23,6 +24,14 @@ func safelyRenameFile(from string, to string) {
 	}
 }
 
+func tryRemoveFile(filename string) {
+	log.Debug("Removing temp file: %v", filename)
+	removeErr := os.Remove(filename)
+	if removeErr != nil {
+		log.Errorf("Failed to remove temp file: %v", filename)
+	}
+}
+
 // thanks https://stackoverflow.com/questions/21060945/simple-way-to-copy-a-file-in-golang
 func copyFileContents(src, dst string) (err error) {
 	// ignore CWE-22 gosec issue - that's more targetted for http based apps that run in a public directory,
@@ -44,6 +53,13 @@ func copyFileContents(src, dst string) (err error) {
 	return out.Sync()
 }
 
+func SafelyCloseReader(reader io.Reader) {
+	switch reader := reader.(type) {
+	case *os.File:
+		safelyCloseFile(reader)
+	}
+}
+
 func safelyCloseFile(file *os.File) {
 	err := file.Close()
 	if err != nil {
@@ -51,3 +67,22 @@ func safelyCloseFile(file *os.File) {
 		log.Error(err.Error())
 	}
 }
+
+func createTempFile() (*os.File, error) {
+	_, err := os.Stat(os.TempDir())
+	if os.IsNotExist(err) {
+		err = os.Mkdir(os.TempDir(), 0700)
+		if err != nil {
+			return nil, err
+		}
+	} else if err != nil {
+		return nil, err
+	}
+
+	file, err := ioutil.TempFile("", "temp")
+	if err != nil {
+		return nil, err
+	}
+
+	return file, err
+}

pkg/yqlib/front_matter.go

@@ -0,0 +1,96 @@
+package yqlib
+
+import (
+	"bufio"
+	"io"
+	"os"
+)
+
+type frontMatterHandler interface {
+	Split() error
+	GetYamlFrontMatterFilename() string
+	GetContentFilename() string
+	CleanUp()
+}
+
+type frontMatterHandlerImpl struct {
+	originalFilename        string
+	yamlFrontMatterFilename string
+	contentFilename         string
+}
+
+func NewFrontMatterHandler(originalFilename string) frontMatterHandler {
+	return &frontMatterHandlerImpl{originalFilename, "", ""}
+}
+
+func (f *frontMatterHandlerImpl) GetYamlFrontMatterFilename() string {
+	return f.yamlFrontMatterFilename
+}
+
+func (f *frontMatterHandlerImpl) GetContentFilename() string {
+	return f.contentFilename
+}
+
+func (f *frontMatterHandlerImpl) CleanUp() {
+	tryRemoveFile(f.yamlFrontMatterFilename)
+	tryRemoveFile(f.contentFilename)
+}
+
+// Splits the given file by yaml front matter
+// yaml content will be saved to first temporary file
+// remaining content will be saved to second temporary file
+func (f *frontMatterHandlerImpl) Split() error {
+	var reader io.Reader
+	var err error
+	if f.originalFilename == "-" {
+		reader = bufio.NewReader(os.Stdin)
+	} else {
+		reader, err = os.Open(f.originalFilename) // #nosec
+		if err != nil {
+			return err
+		}
+	}
+
+	yamlTempFile, err := createTempFile()
+	if err != nil {
+		return err
+	}
+	f.yamlFrontMatterFilename = yamlTempFile.Name()
+	log.Debug("yamlTempFile: %v", yamlTempFile.Name())
+
+	contentTempFile, err := createTempFile()
+	if err != nil {
+		return err
+	}
+	f.contentFilename = contentTempFile.Name()
+	log.Debug("contentTempFile: %v", contentTempFile.Name())
+
+	scanner := bufio.NewScanner(reader)
+
+	lineCount := 0
+	yamlContentBlock := true
+
+	for scanner.Scan() {
+		line := scanner.Text()
+
+		if lineCount > 0 && line == "---" {
+			//we've finished reading the yaml content
+			yamlContentBlock = false
+		}
+		if yamlContentBlock {
+			_, err = yamlTempFile.Write([]byte(line + "\n"))
+		} else {
+			_, err = contentTempFile.Write([]byte(line + "\n"))
+		}
+		if err != nil {
+			return err
+		}
+		lineCount = lineCount + 1
+	}
+
+	safelyCloseFile(yamlTempFile)
+	safelyCloseFile(contentTempFile)
+
+	return scanner.Err()
+
+}

pkg/yqlib/printer.go

@@ -13,8 +13,10 @@ type Printer interface {
 	PrintedAnything() bool
 	SetPrintLeadingSeperator(bool)
 
-	// preamble yaml content
 	SetPreamble(reader io.Reader)
+
+	//e.g. when given a front-matter doc, like jekyll
+	SetAppendix(reader io.Reader)
 }
 
 type resultsPrinter struct {
@@ -30,6 +32,7 @@ type resultsPrinter struct {
 	printedMatches     bool
 	treeNavigator      DataTreeNavigator
 	preambleReader     io.Reader
+	appendixReader     io.Reader
 }
 
 func NewPrinter(writer io.Writer, outputToJSON bool, unwrapScalar bool, colorsEnabled bool, indent int, printDocSeparators bool) Printer {
@@ -56,6 +59,10 @@ func (p *resultsPrinter) SetPreamble(reader io.Reader) {
 	p.preambleReader = reader
 }
 
+func (p *resultsPrinter) SetAppendix(reader io.Reader) {
+	p.appendixReader = reader
+}
+
 func (p *resultsPrinter) PrintedAnything() bool {
 	return p.printedMatches
 }
@@ -139,5 +146,14 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 		p.previousDocIndex = mappedDoc.Document
 	}
 
+	if p.appendixReader != nil && !p.outputToJSON {
+		log.Debug("Piping appendix reader...")
+		betterReader := bufio.NewReader(p.appendixReader)
+		_, err := io.Copy(bufferedWriter, betterReader)
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }

pkg/yqlib/write_in_place_handler.go

@@ -1,7 +1,6 @@
 package yqlib
 
 import (
-	"io/ioutil"
 	"os"
 )
 
@@ -21,27 +20,21 @@ func NewWriteInPlaceHandler(inputFile string) writeInPlaceHandler {
 }
 
 func (w *writeInPlaceHandlerImpl) CreateTempFile() (*os.File, error) {
+	file, err := createTempFile()
+
+	if err != nil {
+		return nil, err
+	}
 	info, err := os.Stat(w.inputFilename)
 	if err != nil {
 		return nil, err
 	}
-	_, err = os.Stat(os.TempDir())
-	if os.IsNotExist(err) {
-		err = os.Mkdir(os.TempDir(), 0700)
-		if err != nil {
-			return nil, err
-		}
-	} else if err != nil {
-		return nil, err
-	}
+	err = os.Chmod(file.Name(), info.Mode())
 
-	file, err := ioutil.TempFile("", "temp")
 	if err != nil {
 		return nil, err
 	}
-
-	err = os.Chmod(file.Name(), info.Mode())
-	log.Debug("writing to tempfile: %v", file.Name())
+	log.Debug("WriteInPlaceHandler: writing to tempfile: %v", file.Name())
 	w.tempFile = file
 	return file, err
 }
@@ -50,13 +43,9 @@ func (w *writeInPlaceHandlerImpl) FinishWriteInPlace(evaluatedSuccessfully bool)
 	log.Debug("Going to write-inplace, evaluatedSuccessfully=%v, target=%v", evaluatedSuccessfully, w.inputFilename)
 	safelyCloseFile(w.tempFile)
 	if evaluatedSuccessfully {
-		log.Debug("moved temp file to target")
+		log.Debug("Moving temp file to target")
 		safelyRenameFile(w.tempFile.Name(), w.inputFilename)
 	} else {
-		log.Debug("removed temp file")
-		removeErr := os.Remove(w.tempFile.Name())
-		if removeErr != nil {
-			log.Errorf("failed removing temp file: %s", w.tempFile.Name())
-		}
+		tryRemoveFile(w.tempFile.Name())
 	}
 }