Array length and collect

cmd/commands_test.go

@@ -190,6 +190,202 @@ func TestReadArrayCmd(t *testing.T) {
 	test.AssertResult(t, "b.e.[1].name: sam\n", result.Output)
 }
 
+func TestReadArrayLengthCmd(t *testing.T) {
+	content := `- things
+- whatever
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l %s", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	test.AssertResult(t, "2\n", result.Output)
+}
+
+func TestReadArrayLengthDeepCmd(t *testing.T) {
+	content := `holder: 
+- things
+- whatever
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l %s holder", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	test.AssertResult(t, "2\n", result.Output)
+}
+
+func TestReadArrayLengthDeepMultipleCmd(t *testing.T) {
+	content := `holderA: 
+- things
+- whatever
+skipMe:
+- yep
+holderB: 
+- other things
+- cool
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l -c %s holder*", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	test.AssertResult(t, "2\n", result.Output)
+}
+
+func TestReadCollectCmd(t *testing.T) {
+	content := `holderA: yep
+skipMe: not me
+holderB: me too
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -c %s holder*", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	expectedOutput := `- yep
+- me too
+`
+	test.AssertResult(t, expectedOutput, result.Output)
+}
+
+func TestReadCollectArrayCmd(t *testing.T) {
+	content := `- name: fred
+  value: 32
+- name: sam
+  value: 67
+- name: fernie
+  value: 103
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -c %s (name==f*)", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	expectedOutput := `- name: fred
+  value: 32
+- name: fernie
+  value: 103
+`
+	test.AssertResult(t, expectedOutput, result.Output)
+}
+
+func TestReadArrayLengthDeepMultipleWithPathCmd(t *testing.T) {
+	content := `holderA: 
+- things
+- whatever
+holderB: 
+- other things
+- cool
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l %s -ppv holder*", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	test.AssertResult(t, "holderA: 2\nholderB: 2\n", result.Output)
+}
+
+func TestReadObjectLengthCmd(t *testing.T) {
+	content := `cat: meow
+dog: bark
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l %s", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	test.AssertResult(t, "2\n", result.Output)
+}
+
+func TestReadObjectLengthDeepCmd(t *testing.T) {
+	content := `holder: 
+  cat: meow
+  dog: bark
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l %s holder", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	test.AssertResult(t, "2\n", result.Output)
+}
+
+func TestReadObjectLengthDeepMultipleCmd(t *testing.T) {
+	content := `holderA: 
+  cat: meow
+  dog: bark
+holderB: 
+  elephant: meow
+  zebra: bark
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l -c %s holder*", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	test.AssertResult(t, "2\n", result.Output)
+}
+
+func TestReadObjectLengthDeepMultipleWithPathsCmd(t *testing.T) {
+	content := `holderA: 
+  cat: meow
+  dog: bark
+holderB: 
+  elephant: meow
+  zebra: bark
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l -ppv %s holder*", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	test.AssertResult(t, "holderA: 2\nholderB: 2\n", result.Output)
+}
+
+func TestReadScalarLengthCmd(t *testing.T) {
+	content := `meow`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l %s", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	test.AssertResult(t, "4\n", result.Output)
+}
+
 func TestReadDeepSplatCmd(t *testing.T) {
 	cmd := getRootCommand()
 	result := test.RunCmd(cmd, "read -p pv ../examples/sample.yaml b.**")

cmd/constant.go

@@ -7,12 +7,15 @@ import (
 
 var customTag = ""
 var printMode = "v"
+var printLength = false
+var collectIntoArray = false
 var writeInplace = false
 var writeScript = ""
 var sourceYamlFile = ""
 var outputToJSON = false
 var prettyPrint = false
 var explodeAnchors = false
+var colorsEnabled = false
 var defaultValue = ""
 var indent = 2
 var overwriteFlag = false

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())
-	encoder.SetIndent(2)
-	errorEncoding := encoder.Encode(&newNode)
-	encoder.Close()
-	return errorEncoding
+	var encoder = yqlib.NewYamlEncoder(cmd.OutOrStdout(), indent, colorsEnabled)
+	return encoder.Encode(&newNode)
 }

cmd/read.go

@@ -26,6 +26,8 @@ yq r -- things.yaml '--key-starting-with-dashes.blah'
 	cmdRead.PersistentFlags().StringVarP(&docIndex, "doc", "d", "0", "process document index number (0 based, * for all documents)")
 	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(&printLength, "length", "l", false, "print length of results")
+	cmdRead.PersistentFlags().BoolVarP(&collectIntoArray, "collect", "c", false, "collect results into array")
 	cmdRead.PersistentFlags().BoolVarP(&explodeAnchors, "explodeAnchors", "X", false, "explode anchors")
 	return cmdRead
 }
@@ -49,6 +51,7 @@ func readProperty(cmd *cobra.Command, args []string) error {
 	if errorReadingStream != nil {
 		return errorReadingStream
 	}
+	out := cmd.OutOrStdout()
 
-	return printResults(matchingNodes, cmd.OutOrStdout())
+	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,12 +77,28 @@ 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
+func lengthOf(node *yaml.Node) int {
+	kindToCheck := node.Kind
+	if node.Kind == yaml.DocumentNode && len(node.Content) == 1 {
+		log.Debugf("length of document node, calculating length of child")
+		kindToCheck = node.Content[0].Kind
 	}
-	return printNode(node, writer)
+	switch kindToCheck {
+	case yaml.ScalarNode:
+		return len(node.Value)
+	case yaml.MappingNode:
+		return len(node.Content) / 2
+	default:
+		return len(node.Content)
+	}
+}
+
+// transforms node before printing, if required
+func transformNode(node *yaml.Node) *yaml.Node {
+	if printLength {
+		return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprintf("%v", lengthOf(node))}
+	}
+	return node
 }
 
 func printNode(node *yaml.Node, writer io.Writer) error {
@@ -90,7 +106,7 @@ func printNode(node *yaml.Node, writer io.Writer) error {
 	if outputToJSON {
 		encoder = yqlib.NewJsonEncoder(writer, prettyPrint, indent)
 	} else {
-		encoder = yqlib.NewYamlEncoder(writer, indent)
+		encoder = yqlib.NewYamlEncoder(writer, indent, colorsEnabled)
 	}
 	return encoder.Encode(node)
 }
@@ -160,6 +176,9 @@ func printResults(matchingNodes []*yqlib.NodeContext, writer io.Writer) error {
 		return nil
 	}
 	var errorWriting error
+
+	var arrayCollection = yaml.Node{Kind: yaml.SequenceNode}
+
 	for _, mappedDoc := range matchingNodes {
 		switch printMode {
 		case "p":
@@ -172,17 +191,27 @@ func printResults(matchingNodes []*yqlib.NodeContext, writer io.Writer) error {
 			var parentNode = yaml.Node{Kind: yaml.MappingNode}
 			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 {
+			parentNode.Content[1] = transformNode(mappedDoc.Node)
+			if collectIntoArray {
+				arrayCollection.Content = append(arrayCollection.Content, &parentNode)
+			} else if err := printNode(&parentNode, bufferedWriter); err != nil {
 				return err
 			}
 		default:
-			if err := printValue(mappedDoc.Node, bufferedWriter); err != nil {
+			if collectIntoArray {
+				arrayCollection.Content = append(arrayCollection.Content, mappedDoc.Node)
+			} else if err := printNode(transformNode(mappedDoc.Node), bufferedWriter); err != nil {
 				return err
 			}
 		}
 	}
 
+	if collectIntoArray {
+		if err := printNode(transformNode(&arrayCollection), bufferedWriter); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -332,8 +361,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))
 }
 

go.mod

@@ -1,6 +1,8 @@
 module github.com/mikefarah/yq/v3
 
 require (
+	github.com/fatih/color v1.9.0
+	github.com/goccy/go-yaml v1.3.2
 	github.com/kr/pretty v0.1.0 // indirect
 	github.com/kylelemons/godebug v1.1.0
 	github.com/pkg/errors v0.8.1

go.sum

@@ -4,8 +4,16 @@ github.com/coreos/etcd v3.3.10+incompatible/go.mod h1:uF7uidLiAD3TWHmW31ZFd/JWoc
 github.com/coreos/go-etcd v2.0.0+incompatible/go.mod h1:Jez6KQU2B/sWsbdaef3ED8NzMklzPG4d5KIOhIy30Tk=
 github.com/coreos/go-semver v0.2.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/cpuguy83/go-md2man v1.0.10/go.mod h1:SmD6nW6nTyfqj6ABTjUi3V3JVMnlJmwcJI5acqYI6dE=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
+github.com/fatih/color v1.9.0 h1:8xPHl4/q1VyqGIPif1F+1V3Y3lSmrq01EabUW3CoW5s=
+github.com/fatih/color v1.9.0/go.mod h1:eQcE1qtQxscV5RaZvpXrrb8Drkc3/DdQ+uUYCNjL+zU=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
+github.com/go-playground/locales v0.13.0/go.mod h1:taPMhCMXrRLJO55olJkUXHZBHCxTMfnGwq/HNwmWNS8=
+github.com/go-playground/universal-translator v0.17.0/go.mod h1:UkSxE5sNxxRwHyU+Scu5vgOQjsIJAF8j9muTVoKLVtA=
+github.com/goccy/go-yaml v1.3.2 h1:joykVKVARE+kQNoaj0ijjPY7lhgdovyU6etuYEl3hFU=
+github.com/goccy/go-yaml v1.3.2/go.mod h1:PsEEJ29nIFZL07P/c8dv4P6rQkVFFXafQee85U+ERHA=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
 github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
@@ -16,7 +24,14 @@ github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
 github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
+github.com/leodido/go-urn v1.2.0/go.mod h1:+8+nEpDfqqsY+g338gtMEUOtuK+4dEMhiQEgxpxOKII=
 github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
+github.com/mattn/go-colorable v0.1.4 h1:snbPLB8fVfU9iwbbo30TPtbLRzwWu6aJS6Xh4eaaviA=
+github.com/mattn/go-colorable v0.1.4/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
+github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
+github.com/mattn/go-isatty v0.0.10/go.mod h1:qgIWMr58cqv1PHHyhnkY9lrL7etaEgOFcMEpPG5Rm84=
+github.com/mattn/go-isatty v0.0.11 h1:FxPOTFNqGkuDUGi3H/qkUbQO4ZiBa2brKq5r0l8TGeM=
+github.com/mattn/go-isatty v0.0.11/go.mod h1:PhnuNfih5lzO57/f3n+odYbM4JtupLOxQOAqxQCu2WE=
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
 github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
@@ -34,19 +49,33 @@ github.com/spf13/pflag v1.0.3/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnIn
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/spf13/viper v1.3.2/go.mod h1:ZiWeW+zYFKm7srdB9IoDzzZXaJaI5eL9QjNiN/DMA2s=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
+github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/ugorji/go/codec v0.0.0-20181204163529-d75b2dcb6bc8/go.mod h1:VFNgLljTbGfSG7qAOspJ7OScBnGdDN/yBr0sguwnwf0=
 github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
 golang.org/x/crypto v0.0.0-20181203042331-505ab145d0a9/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
 golang.org/x/sys v0.0.0-20181205085412-a5c9d58dba9a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20191008105621-543471e840be/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20191010194322-b09406accb47/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20191026070338-33540a1f6037 h1:YyJpGZS1sBuBCzLAR1VEpK193GlqGZbnPFnPV/5Rsb4=
+golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898 h1:/atklqdjdhuosWIl6AIbOeHJjicWYPqR9bpxqxYG2pA=
+golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/go-playground/assert.v1 v1.2.1/go.mod h1:9RXL0bg/zibRAgZUYszZSwO/z8Y/a8bDuhia5mkpMnE=
+gopkg.in/go-playground/validator.v9 v9.30.0/go.mod h1:+c9/zcJMFNgbLvly1L1V+PpxWdVbfP1avr/N00E2vyQ=
 gopkg.in/op/go-logging.v1 v1.0.0-20160211212156-b2cb9fa56473 h1:6D+BvnJ/j6e222UW8s2qTSe3wGBtvo0MbVQG/c5k8RE=
 gopkg.in/op/go-logging.v1 v1.0.0-20160211212156-b2cb9fa56473/go.mod h1:N1eN2tsCx0Ydtgjl4cqmbRCsY4/+z4cYDeqwZTk6zog=
 gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v3 v3.0.0-20191120175047-4206685974f2/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.0-20200121175148-a6ecf24a6d71 h1:Xe2gvTZUJpsvOWUnvmL/tmhVBZUmHSvLbMjRj6NUUKo=
 gopkg.in/yaml.v3 v3.0.0-20200121175148-a6ecf24a6d71/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

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 {
-	var encoder = yaml.NewEncoder(destination)
+func NewYamlEncoder(destination io.Writer, indent int, colorise bool) Encoder {
 	if indent < 0 {
 		indent = 0
 	}
-	encoder.SetIndent(indent)
-	return &yamlEncoder{encoder}
+	return &yamlEncoder{destination, indent, colorise, true}
 }
 
 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