diff --git a/cmd/new.go b/cmd/new.go index 3ca842d..a57f2ed 100644 --- a/cmd/new.go +++ b/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) } diff --git a/cmd/read.go b/cmd/read.go index 7946abe..4a30faa 100644 --- a/cmd/read.go +++ b/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) } diff --git a/cmd/root.go b/cmd/root.go index 844fde1..24d8b2e 100644 --- a/cmd/root.go +++ b/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(), diff --git a/cmd/utils.go b/cmd/utils.go index 257edd9..93a674a 100644 --- a/cmd/utils.go +++ b/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)) } diff --git a/pkg/yqlib/color_print.go b/pkg/yqlib/color_print.go new file mode 100644 index 0000000..82968ed --- /dev/null +++ b/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 +} diff --git a/pkg/yqlib/encoder.go b/pkg/yqlib/encoder.go index 5a4be04..172fa38 100644 --- a/pkg/yqlib/encoder.go +++ b/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 { diff --git a/scripts/format.sh b/scripts/format.sh index 2acfd35..123155f 100755 --- a/scripts/format.sh +++ b/scripts/format.sh @@ -1,3 +1,4 @@ #!/bin/bash find . \( -path ./vendor \) -prune -o -name "*.go" -exec goimports -w {} \; +go mod tidy