Colors work for all commands

cmd/new.go

@@ -1,8 +1,8 @@
 package cmd
 
 import (
+	"github.com/mikefarah/yq/v3/pkg/yqlib"
 	"github.com/spf13/cobra"
-	yaml "gopkg.in/yaml.v3"
 )
 
 func createNewCmd() *cobra.Command {
@@ -46,9 +46,6 @@ func newProperty(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	var encoder = yaml.NewEncoder(cmd.OutOrStdout())
+	var encoder = yqlib.NewYamlEncoder(cmd.OutOrStdout(), indent, colorsEnabled)
-	encoder.SetIndent(2)
+	return encoder.Encode(&newNode)
-	errorEncoding := encoder.Encode(&newNode)
-	encoder.Close()
-	return errorEncoding
 }

cmd/read.go

@@ -1,66 +1,10 @@
 package cmd
 
 import (
-	"bytes"
-	"fmt"
-	"io"
-
-	"github.com/fatih/color"
-	"github.com/goccy/go-yaml/lexer"
-	"github.com/goccy/go-yaml/printer"
 	errors "github.com/pkg/errors"
 	"github.com/spf13/cobra"
 )
 
-const escape = "\x1b"
-
-func format(attr color.Attribute) string {
-	return fmt.Sprintf("%s[%dm", escape, attr)
-}
-
-func colorizeAndPrint(bytes []byte, writer io.Writer) error {
-	tokens := lexer.Tokenize(string(bytes))
-	var p printer.Printer
-	p.Bool = func() *printer.Property {
-		return &printer.Property{
-			Prefix: format(color.FgHiMagenta),
-			Suffix: format(color.Reset),
-		}
-	}
-	p.Number = func() *printer.Property {
-		return &printer.Property{
-			Prefix: format(color.FgHiMagenta),
-			Suffix: format(color.Reset),
-		}
-	}
-	p.MapKey = func() *printer.Property {
-		return &printer.Property{
-			Prefix: format(color.FgHiCyan),
-			Suffix: format(color.Reset),
-		}
-	}
-	p.Anchor = func() *printer.Property {
-		return &printer.Property{
-			Prefix: format(color.FgHiYellow),
-			Suffix: format(color.Reset),
-		}
-	}
-	p.Alias = func() *printer.Property {
-		return &printer.Property{
-			Prefix: format(color.FgHiYellow),
-			Suffix: format(color.Reset),
-		}
-	}
-	p.String = func() *printer.Property {
-		return &printer.Property{
-			Prefix: format(color.FgHiGreen),
-			Suffix: format(color.Reset),
-		}
-	}
-	_, err := writer.Write([]byte(p.PrintTokens(tokens) + "\n"))
-	return err
-}
-
 func createReadCmd() *cobra.Command {
 	var cmdRead = &cobra.Command{
 		Use:     "read [yaml_file] [path_expression]",
@@ -83,7 +27,6 @@ yq r -- things.yaml '--key-starting-with-dashes.blah'
 	cmdRead.PersistentFlags().StringVarP(&printMode, "printMode", "p", "v", "print mode (v (values, default), p (paths), pv (path and value pairs)")
 	cmdRead.PersistentFlags().StringVarP(&defaultValue, "defaultValue", "D", "", "default value printed when there are no results")
 	cmdRead.PersistentFlags().BoolVarP(&explodeAnchors, "explodeAnchors", "X", false, "explode anchors")
-	cmdRead.PersistentFlags().BoolVarP(&colorsEnabled, "colorsEnabled", "C", false, "enable colors")
 	return cmdRead
 }
 
@@ -108,12 +51,5 @@ func readProperty(cmd *cobra.Command, args []string) error {
 	}
 	out := cmd.OutOrStdout()
 
-	if colorsEnabled && !outputToJSON {
-		buffer := bytes.NewBuffer(nil)
-		if err := printResults(matchingNodes, buffer); err != nil {
-			return err
-		}
-		return colorizeAndPrint(buffer.Bytes(), out)
-	}
 	return printResults(matchingNodes, out)
 }

cmd/root.go

@@ -43,6 +43,7 @@ func New() *cobra.Command {
 	rootCmd.PersistentFlags().BoolVarP(&prettyPrint, "prettyPrint", "P", false, "pretty print")
 	rootCmd.PersistentFlags().IntVarP(&indent, "indent", "I", 2, "sets indent level for output")
 	rootCmd.Flags().BoolVarP(&version, "version", "V", false, "Print version information and quit")
+	rootCmd.PersistentFlags().BoolVarP(&colorsEnabled, "colorsEnabled", "C", false, "enable colors")
 
 	rootCmd.AddCommand(
 		createReadCmd(),

cmd/utils.go

@@ -77,20 +77,12 @@ func appendDocument(originalMatchingNodes []*yqlib.NodeContext, dataBucket yaml.
 	return append(originalMatchingNodes, matchingNodes...), nil
 }
 
-func printValue(node *yaml.Node, writer io.Writer) error {
-	if node.Kind == yaml.ScalarNode {
-		_, errorWriting := writer.Write([]byte(node.Value + "\n"))
-		return errorWriting
-	}
-	return printNode(node, writer)
-}
-
 func printNode(node *yaml.Node, writer io.Writer) error {
 	var encoder yqlib.Encoder
 	if outputToJSON {
 		encoder = yqlib.NewJsonEncoder(writer, prettyPrint, indent)
 	} else {
-		encoder = yqlib.NewYamlEncoder(writer, indent)
+		encoder = yqlib.NewYamlEncoder(writer, indent, colorsEnabled)
 	}
 	return encoder.Encode(node)
 }
@@ -173,11 +165,11 @@ func printResults(matchingNodes []*yqlib.NodeContext, writer io.Writer) error {
 			parentNode.Content = make([]*yaml.Node, 2)
 			parentNode.Content[0] = &yaml.Node{Kind: yaml.ScalarNode, Value: lib.PathStackToString(mappedDoc.PathStack)}
 			parentNode.Content[1] = mappedDoc.Node
-			if err := printValue(&parentNode, bufferedWriter); err != nil {
+			if err := printNode(&parentNode, bufferedWriter); err != nil {
 				return err
 			}
 		default:
-			if err := printValue(mappedDoc.Node, bufferedWriter); err != nil {
+			if err := printNode(mappedDoc.Node, bufferedWriter); err != nil {
 				return err
 			}
 		}
@@ -332,8 +324,9 @@ func readAndUpdate(stdOut io.Writer, inputFile string, updateData updateDataFn)
 	if outputToJSON {
 		encoder = yqlib.NewJsonEncoder(bufferedWriter, prettyPrint, indent)
 	} else {
-		encoder = yqlib.NewYamlEncoder(bufferedWriter, indent)
+		encoder = yqlib.NewYamlEncoder(bufferedWriter, indent, colorsEnabled)
 	}
+
 	return readStream(inputFile, mapYamlDecoder(updateData, encoder))
 }
 

pkg/yqlib/color_print.go

@@ -0,0 +1,61 @@
+package yqlib
+
+import (
+	"fmt"
+	"io"
+
+	"github.com/fatih/color"
+	"github.com/goccy/go-yaml/lexer"
+	"github.com/goccy/go-yaml/printer"
+)
+
+// Thanks @risentveber!
+
+const escape = "\x1b"
+
+func format(attr color.Attribute) string {
+	return fmt.Sprintf("%s[%dm", escape, attr)
+}
+
+func ColorizeAndPrint(bytes []byte, writer io.Writer) error {
+	tokens := lexer.Tokenize(string(bytes))
+	var p printer.Printer
+	p.Bool = func() *printer.Property {
+		return &printer.Property{
+			Prefix: format(color.FgHiMagenta),
+			Suffix: format(color.Reset),
+		}
+	}
+	p.Number = func() *printer.Property {
+		return &printer.Property{
+			Prefix: format(color.FgHiMagenta),
+			Suffix: format(color.Reset),
+		}
+	}
+	p.MapKey = func() *printer.Property {
+		return &printer.Property{
+			Prefix: format(color.FgCyan),
+			Suffix: format(color.Reset),
+		}
+	}
+	p.Anchor = func() *printer.Property {
+		return &printer.Property{
+			Prefix: format(color.FgHiYellow),
+			Suffix: format(color.Reset),
+		}
+	}
+	p.Alias = func() *printer.Property {
+		return &printer.Property{
+			Prefix: format(color.FgHiYellow),
+			Suffix: format(color.Reset),
+		}
+	}
+	p.String = func() *printer.Property {
+		return &printer.Property{
+			Prefix: format(color.FgGreen),
+			Suffix: format(color.Reset),
+		}
+	}
+	_, err := writer.Write([]byte(p.PrintTokens(tokens) + "\n"))
+	return err
+}

pkg/yqlib/encoder.go

@@ -1,6 +1,7 @@
 package yqlib
 
 import (
+	"bytes"
 	"encoding/json"
 	"io"
 
@@ -12,20 +13,45 @@ type Encoder interface {
 }
 
 type yamlEncoder struct {
-	encoder *yaml.Encoder
+	destination io.Writer
+	indent      int
+	colorise    bool
+	firstDoc    bool
 }
 
-func NewYamlEncoder(destination io.Writer, indent int) Encoder {
+func NewYamlEncoder(destination io.Writer, indent int, colorise bool) Encoder {
-	var encoder = yaml.NewEncoder(destination)
 	if indent < 0 {
 		indent = 0
 	}
-	encoder.SetIndent(indent)
+	return &yamlEncoder{destination, indent, colorise, true}
-	return &yamlEncoder{encoder}
 }
 
 func (ye *yamlEncoder) Encode(node *yaml.Node) error {
-	return ye.encoder.Encode(node)
+
+	destination := ye.destination
+	tempBuffer := bytes.NewBuffer(nil)
+	if ye.colorise {
+		destination = tempBuffer
+	}
+
+	var encoder = yaml.NewEncoder(destination)
+
+	encoder.SetIndent(ye.indent)
+	// TODO: work out if the first doc had a separator or not.
+	if ye.firstDoc {
+		ye.firstDoc = false
+	} else if _, err := destination.Write([]byte("---\n")); err != nil {
+		return err
+	}
+
+	if err := encoder.Encode(node); err != nil {
+		return err
+	}
+
+	if ye.colorise {
+		return ColorizeAndPrint(tempBuffer.Bytes(), ye.destination)
+	}
+	return nil
 }
 
 type jsonEncoder struct {

scripts/format.sh

@@ -1,3 +1,4 @@
 #!/bin/bash
 
 find . \( -path ./vendor \) -prune -o -name "*.go" -exec goimports -w {} \;
+go mod tidy