From 297522cbddfcb62fb6514ff7137853657dfd8242 Mon Sep 17 00:00:00 2001 From: Mike Farah Date: Wed, 13 Jun 2018 14:10:00 +1000 Subject: [PATCH] Write supports multidoc yaml, better use of yaml library streaming --- commands_test.go | 24 +++++- vendor/vendor.json | 8 +- yq.go | 182 +++++++++++++++++++++++---------------------- yq_test.go | 31 -------- 4 files changed, 122 insertions(+), 123 deletions(-) diff --git a/commands_test.go b/commands_test.go index 86d2e068..a72f9542 100644 --- a/commands_test.go +++ b/commands_test.go @@ -372,6 +372,28 @@ func TestWriteCmd(t *testing.T) { assertResult(t, expectedOutput, result.Output) } +func TestWriteMultiCmd(t *testing.T) { + content := `b: + c: 3 +--- +apples: great +` + filename := writeTempYamlFile(content) + defer removeTempYamlFile(filename) + + cmd := getRootCommand() + result := runCmd(cmd, fmt.Sprintf("write %s -d 1 apples ok", filename)) + if result.Error != nil { + t.Error(result.Error) + } + expectedOutput := `b: + c: 3 +--- +apples: ok +` + assertResult(t, expectedOutput, result.Output) +} + func TestWriteCmd_EmptyArray(t *testing.T) { content := `b: 3` filename := writeTempYamlFile(content) @@ -441,7 +463,7 @@ func TestWriteCmd_Inplace(t *testing.T) { gotOutput := readTempYamlFile(filename) expectedOutput := `b: c: 7` - assertResult(t, expectedOutput, gotOutput) + assertResult(t, expectedOutput, strings.Trim(gotOutput, "\n ")) } func TestWriteCmd_Append(t *testing.T) { diff --git a/vendor/vendor.json b/vendor/vendor.json index 92fd549a..5529e16f 100644 --- a/vendor/vendor.json +++ b/vendor/vendor.json @@ -33,10 +33,12 @@ "revisionTime": "2017-08-24T17:57:12Z" }, { - "checksumSHA1": "ZSWoOPUNRr5+3dhkLK3C4cZAQPk=", + "checksumSHA1": "DHNYKS5T54/XOqUsFFzdZMLEnVE=", + "origin": "github.com/mikefarah/yaml", "path": "gopkg.in/yaml.v2", - "revision": "5420a8b6744d3b0345ab293f6fcba19c978f1183", - "revisionTime": "2018-03-28T19:50:20Z" + "revision": "e175af14aaa1d0eff2ee04b691e4a4827a111416", + "revisionTime": "2018-06-13T04:05:11Z", + "tree": true } ], "rootPath": "github.com/mikefarah/yq" diff --git a/yq.go b/yq.go index 14348456..c03e2a6f 100644 --- a/yq.go +++ b/yq.go @@ -7,6 +7,7 @@ import ( "io" "io/ioutil" "os" + "reflect" "strconv" "strings" @@ -34,6 +35,7 @@ func main() { } func newCommandCLI() *cobra.Command { + yaml.DefaultMapType = reflect.TypeOf(yaml.MapSlice{}) var rootCmd = &cobra.Command{ Use: "yq", RunE: func(cmd *cobra.Command, args []string) error { @@ -102,15 +104,15 @@ func createWriteCmd() *cobra.Command { var cmdWrite = &cobra.Command{ Use: "write [yaml_file] [path] [value]", Aliases: []string{"w"}, - Short: "yq w [--inplace/-i] [--script/-s script_file] sample.yaml a.b.c newValueForC", + Short: "yq w [--inplace/-i] [--script/-s script_file] [--doc/-d document_index] sample.yaml a.b.c newValueForC", Example: ` yq write things.yaml a.b.c cat yq write --inplace things.yaml a.b.c cat yq w -i things.yaml a.b.c cat yq w --script update_script.yaml things.yaml yq w -i -s update_script.yaml things.yaml -yq w things.yaml a.b.d[+] foo -yq w things.yaml a.b.d[+] foo +yq w --doc 2 things.yaml a.b.d[+] foo +yq w -d2 things.yaml a.b.d[+] foo `, Long: `Updates the yaml file w.r.t the given path and value. Outputs to STDOUT unless the inplace flag is used, in which case the file is updated instead. @@ -129,6 +131,7 @@ a.b.e: } cmdWrite.PersistentFlags().BoolVarP(&writeInplace, "inplace", "i", false, "update the yaml file inplace") cmdWrite.PersistentFlags().StringVarP(&writeScript, "script", "s", "", "yaml script for updating yaml") + cmdWrite.PersistentFlags().IntVarP(&docIndex, "doc", "d", 0, "process document index number (0 based)") return cmdWrite } @@ -212,7 +215,6 @@ func readProperty(cmd *cobra.Command, args []string) error { } func read(args []string) (interface{}, error) { - var parsedData yaml.MapSlice var path = "" if len(args) < 1 { @@ -220,67 +222,17 @@ func read(args []string) (interface{}, error) { } else if len(args) > 1 { path = args[1] } - - if err := readData(args[0], docIndex, &parsedData); err != nil { - var generalData interface{} - if err = readData(args[0], docIndex, &generalData); err != nil { - return nil, err - } - item := yaml.MapItem{Key: "thing", Value: generalData} - parsedData = yaml.MapSlice{item} - path = "thing." + path + var generalData interface{} + if err := readData(args[0], docIndex, &generalData); err != nil { + return nil, err } - - if parsedData != nil && parsedData[0].Key == nil { - var interfaceData []map[interface{}]interface{} - if err := readData(args[0], docIndex, &interfaceData); err == nil { - var listMap []yaml.MapSlice - for _, item := range interfaceData { - listMap = append(listMap, mapToMapSlice(item)) - } - return readYamlArray(listMap, path) - } - } - if path == "" { - return parsedData, nil + return generalData, nil } var paths = parsePath(path) - - return readMap(parsedData, paths[0], paths[1:]) -} - -func readYamlArray(listMap []yaml.MapSlice, path string) (interface{}, error) { - if path == "" { - return listMap, nil - } - - var paths = parsePath(path) - - if paths[0] == "*" { - if len(paths[1:]) == 0 { - return listMap, nil - } - var results []interface{} - for _, m := range listMap { - value, err := readMap(m, paths[1], paths[2:]) - if err != nil { - return nil, err - } - results = append(results, value) - } - return results, nil - } - - index, err := strconv.ParseInt(paths[0], 10, 64) - if err != nil { - return nil, fmt.Errorf("Error accessing array: %v", err) - } - if len(paths[1:]) == 0 { - return listMap[index], nil - } - return readMap(listMap[index], paths[1], paths[2:]) + value, err := recurse(generalData, paths[0], paths[1:]) + return value, err } func newProperty(cmd *cobra.Command, args []string) error { @@ -316,12 +268,73 @@ func newYaml(args []string) (interface{}, error) { return updateParsedData(parsedData, writeCommands, prependCommand) } -func writeProperty(cmd *cobra.Command, args []string) error { - updatedData, err := updateYaml(args) - if err != nil { - return err +func mapYamlDecoder(writeCommands yaml.MapSlice, encoder *yaml.Encoder) yamlDecoderFn { + return func(decoder *yaml.Decoder) error { + var dataBucket interface{} + var errorReading error + var errorWriting error + var currentIndex = 0 + + for { + log.Debugf("Read doc %v", currentIndex) + errorReading = decoder.Decode(&dataBucket) + + if errorReading == io.EOF { + return nil + } else if errorReading != nil { + return fmt.Errorf("Error reading document at index %v, %v", currentIndex, errorReading) + } + + if currentIndex == docIndex { + log.Debugf("Updating doc %v", currentIndex) + for _, entry := range writeCommands { + path := entry.Key.(string) + value := entry.Value + log.Debugf("setting %v to %v", path, value) + var paths = parsePath(path) + dataBucket = updatedChildValue(dataBucket, paths, value) + } + } + + errorWriting = encoder.Encode(dataBucket) + + if errorWriting != nil { + return fmt.Errorf("Error writing document at index %v, %v", currentIndex, errorWriting) + } + currentIndex = currentIndex + 1 + } } - return write(cmd, args[0], updatedData) +} + +func writeProperty(cmd *cobra.Command, args []string) error { + var writeCommands, writeCommandsError = readWriteCommands(args, 3, "Must provide ") + if writeCommandsError != nil { + return writeCommandsError + } + var inputFile = args[0] + var destination io.Writer + var destinationName string + if writeInplace { + var tempFile, err = ioutil.TempFile("", "temp") + if err != nil { + return err + } + destinationName = tempFile.Name() + destination = tempFile + defer func() { + safelyCloseFile(tempFile) + safelyRenameFile(tempFile.Name(), inputFile) + }() + } else { + var writer = bufio.NewWriter(cmd.OutOrStdout()) + destination = writer + destinationName = "Stdout" + defer safelyFlush(writer) + } + var encoder = yaml.NewEncoder(destination) + log.Debugf("Writing to %v from %v", destinationName, inputFile) + //need to use a temp file if writeInplace is given + return readStream(inputFile, mapYamlDecoder(writeCommands, encoder)) } func write(cmd *cobra.Command, filename string, updatedData interface{}) error { @@ -432,27 +445,6 @@ func readWriteCommands(args []string, expectedArgs int, badArgsMessage string) ( return writeCommands, nil } -func updateYaml(args []string) (interface{}, error) { - var writeCommands, writeCommandsError = readWriteCommands(args, 3, "Must provide ") - if writeCommandsError != nil { - return nil, writeCommandsError - } - - var prependCommand = "" - var parsedData yaml.MapSlice - if err := readData(args[0], 0, &parsedData); err != nil { - var generalData interface{} - if err = readData(args[0], 0, &generalData); err != nil { - return nil, err - } - item := yaml.MapItem{Key: "thing", Value: generalData} - parsedData = yaml.MapSlice{item} - prependCommand = "thing" - } - - return updateParsedData(parsedData, writeCommands, prependCommand) -} - func parseValue(argument string) interface{} { var value, err interface{} var inQuotes = len(argument) > 0 && argument[0] == '"' @@ -505,11 +497,25 @@ func marshalContext(context interface{}) (string, error) { return outStr, nil } +func safelyRenameFile(from string, to string) { + if err := os.Rename(from, to); err != nil { + log.Errorf("Error renaming from %v to %v", from, to) + log.Error(err.Error()) + } +} + +func safelyFlush(writer *bufio.Writer) { + if err := writer.Flush(); err != nil { + log.Error("Error flushing writer!") + log.Error(err.Error()) + } + +} func safelyCloseFile(file *os.File) { err := file.Close() if err != nil { - fmt.Println("Error closing file!") - fmt.Println(err.Error()) + log.Error("Error closing file!") + log.Error(err.Error()) } } diff --git a/yq_test.go b/yq_test.go index 5eda4fc1..ff8cf5c8 100644 --- a/yq_test.go +++ b/yq_test.go @@ -78,37 +78,6 @@ func TestNewYamlArray(t *testing.T) { formattedResult) } -func TestUpdateYaml(t *testing.T) { - result, _ := updateYaml([]string{"examples/sample.yaml", "b.c", "3"}) - formattedResult := fmt.Sprintf("%v", result) - assertResult(t, - "[{a Easy! as one two three} {b [{c 3} {d [3 4]} {e [[{name fred} {value 3}] [{name sam} {value 4}]]}]}]", - formattedResult) -} - -func TestUpdateYamlArray(t *testing.T) { - result, _ := updateYaml([]string{"examples/sample_array.yaml", "[0]", "3"}) - formattedResult := fmt.Sprintf("%v", result) - assertResult(t, - "[3 2 3]", - formattedResult) -} - -func TestUpdateYaml_WithScript(t *testing.T) { - writeScript = "examples/instruction_sample.yaml" - _, _ = updateYaml([]string{"examples/sample.yaml"}) -} - -func TestUpdateYaml_WithUnknownScript(t *testing.T) { - writeScript = "fake-unknown" - _, err := updateYaml([]string{"examples/sample.yaml"}) - if err == nil { - t.Error("Expected error due to unknown file") - } - expectedOutput := `open fake-unknown: no such file or directory` - assertResult(t, expectedOutput, err.Error()) -} - func TestNewYaml_WithScript(t *testing.T) { writeScript = "examples/instruction_sample.yaml" expectedResult := `b: