update jsonnet with format

drone/jsonnet/jsonnet.go

@@ -1,12 +1,14 @@
 package jsonnet
 
 import (
-	"fmt"
+	"bytes"
+	"io"
 	"io/ioutil"
 	"log"
 	"os"
 
-	"github.com/drone/drone-cli/drone/jsonnet/stdlib"
+	"github.com/drone/drone-yaml/yaml"
+	"github.com/drone/drone-yaml/yaml/pretty"
 	"github.com/fatih/color"
 	"github.com/google/go-jsonnet"
 	"github.com/urfave/cli"
@@ -23,122 +25,89 @@ var Command = cli.Command{
 		}
 	},
 	Flags: []cli.Flag{
+		cli.StringFlag{
+			Name:  "source",
+			Usage: "Source file",
+			Value: ".drone.jsonnet",
+		},
+		cli.StringFlag{
+			Name:  "target",
+			Usage: "target file",
+			Value: ".drone.yml",
+		},
 		cli.BoolFlag{
-			Name: "stream",
+			Name:  "stream",
+			Usage: "Write output as a YAML stream.",
 		},
 		cli.BoolTFlag{
-			Name: "string",
-		},
-		cli.IntFlag{
-			Name:  "max-stack",
-			Usage: "number of allowed stack frames",
-			Value: 500,
-		},
-		cli.IntFlag{
-			Name:  "max-trace",
-			Usage: "max length of stack trace before cropping",
-			Value: 20,
+			Name:  "format",
+			Usage: "Write output as formatted YAML",
 		},
 		cli.BoolFlag{
 			Name:  "stdout",
-			Usage: "write the json document to stdout",
+			Usage: "Write output to stdout",
+		},
+		cli.BoolFlag{
+			Name:  "string",
+			Usage: "Expect a string, manifest as plain text",
 		},
 	},
 }
 
 func generate(c *cli.Context) error {
-	input := c.Args().Get(0)
-	if input == "" {
-		input = ".drone.jsonnet"
-	}
-	output := c.Args().Get(1)
-	if output == "" {
-		output = ".drone.yml"
-	}
+	source := c.String("source")
+	target := c.String("target")
 
-	snippet, err := ioutil.ReadFile(input)
+	data, err := ioutil.ReadFile(source)
 	if err != nil {
 		return err
 	}
 
 	vm := jsonnet.MakeVM()
-	vm.MaxStack = c.Int("max-stack")
-	vm.StringOutput = c.BoolT("string")
-	vm.ErrorFormatter.SetMaxStackTraceSize(
-		c.Int("max-trace"),
-	)
+	vm.MaxStack = 500
+	vm.StringOutput = c.Bool("string")
+	vm.ErrorFormatter.SetMaxStackTraceSize(20)
 	vm.ErrorFormatter.SetColorFormatter(
 		color.New(color.FgRed).Fprintf,
 	)
-	vm.Importer(
-		stdlib.Importer(),
-	)
 
-	var result string
-	var resultArray []string
+	buf := new(bytes.Buffer)
 	if c.Bool("stream") {
-		resultArray, err = vm.EvaluateSnippetStream(input, string(snippet))
+		docs, err := vm.EvaluateSnippetStream(source, string(data))
+		if err != nil {
+			return err
+		}
+		for _, doc := range docs {
+			buf.WriteString("---")
+			buf.WriteString("\n")
+			buf.WriteString(doc)
+		}
 	} else {
-		result, err = vm.EvaluateSnippet(input, string(snippet))
-	}
-	if err != nil {
-		return err
-	}
-
-	if c.Bool("stream") {
-		return writeOutputStream(resultArray, output)
-	}
-	return writeOutputFile(result, output)
-}
-
-// writeOutputStream writes the output as a YAML stream.
-func writeOutputStream(output []string, outputFile string) error {
-	var f *os.File
-
-	if outputFile == "" {
-		f = os.Stdout
-	} else {
-		var err error
-		f, err = os.Create(outputFile)
+		result, err := vm.EvaluateSnippet(source, string(data))
 		if err != nil {
 			return err
 		}
-		defer f.Close()
+		buf.WriteString(result)
 	}
 
-	for _, doc := range output {
-		_, err := f.WriteString("---\n")
-		if err != nil {
-			return err
-		}
-		_, err = f.WriteString(doc)
+	// the yaml file is parsed and formatted by default. This
+	// can be disabled for --format=false.
+	if c.BoolT("format") {
+		manifest, err := yaml.Parse(buf)
 		if err != nil {
 			return err
 		}
+		buf.Reset()
+		pretty.Print(buf, manifest)
 	}
 
-	if len(output) > 0 {
-		_, err := f.WriteString("...\n")
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func writeOutputFile(output string, outputFile string) error {
-	if outputFile == "" {
-		fmt.Print(output)
+	// the user can optionally write the yaml to stdout. This
+	// is useful for debugging purposes without mutating an
+	// existing file.
+	if c.Bool("stdout") {
+		io.Copy(os.Stdout, buf)
 		return nil
 	}
 
-	f, err := os.Create(outputFile)
-	if err != nil {
-		return err
-	}
-	defer f.Close()
-
-	_, err = f.WriteString(output)
-	return err
+	return ioutil.WriteFile(target, buf.Bytes(), 0644)
 }