add colorization

cmd/constant.go

@@ -13,6 +13,7 @@ var sourceYamlFile = ""
 var outputToJSON = false
 var prettyPrint = false
 var explodeAnchors = false
+var colorsEnabled = false
 var defaultValue = ""
 var indent = 2
 var overwriteFlag = false

cmd/read.go

@@ -1,10 +1,66 @@
 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]",
@@ -27,6 +83,7 @@ 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
 }
 
@@ -49,6 +106,14 @@ func readProperty(cmd *cobra.Command, args []string) error {
 	if errorReadingStream != nil {
 		return errorReadingStream
 	}
+	out := cmd.OutOrStdout()
 
-	return printResults(matchingNodes, 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)
 }