cmd/cue/cmd: allow capturing output

Issue #50

Change-Id: I06be2f1e4dc4ce1d531ddb9c35eb3c11c5693f8d
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/2162
Reviewed-by: Marcel van Lohuizen <mpvl@google.com>

cmd/cue/cmd/cmd_test.go

@@ -36,7 +36,7 @@ func TestCmd(t *testing.T) {
 		stderr = os.Stderr
 	}()
 	for _, name := range testCases {
-		rootCmd := newRootCmd()
+		rootCmd := newRootCmd().root
 		run := func(cmd *cobra.Command, args []string) error {
 			stdout = cmd.OutOrStdout()
 			stderr = cmd.OutOrStderr()

cmd/cue/cmd/common.go

@@ -29,40 +29,42 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func init() {
-	s, err := os.Getwd()
-	if err == nil {
-		cwd = s
-	}
-}
-
 var runtime = &cue.Runtime{}
 
-var cwd = "////"
-
 // printHeader is a hacky and unprincipled way to sanatize the package path.
-func printHeader(w io.Writer, dir string) {
-	head := strings.Replace(dir, cwd, ".", 1)
-	fmt.Fprintf(w, "--- %s\n", head)
+func printHeader(w io.Writer, cwd, dir string) {
+	if cwd != "" {
+		if dir == cwd {
+			return
+		}
+		dir = strings.Replace(dir, cwd, ".", 1)
+	}
+	fmt.Fprintf(w, "--- %s\n", dir)
 }
 
 func exitIfErr(cmd *cobra.Command, inst *cue.Instance, err error, fatal bool) {
 	exitOnErr(cmd, inst.Dir, err, fatal)
 }
 
 func exitOnErr(cmd *cobra.Command, file string, err error, fatal bool) {
-	if err != nil {
-		w := &bytes.Buffer{}
-		printHeader(w, file)
-		errors.Print(w, err)
-
-		// TODO: do something more principled than this.
-		b := w.Bytes()
-		b = bytes.ReplaceAll(b, []byte(cwd), []byte("."))
-		cmd.OutOrStderr().Write(b)
-		if fatal {
-			exit()
-		}
+	if err == nil {
+		return
+	}
+	cwd := "////"
+	if p, _ := os.Getwd(); p != "" {
+		cwd = p
+	}
+
+	w := &bytes.Buffer{}
+	printHeader(w, cwd, file)
+	errors.Print(w, err)
+
+	// TODO: do something more principled than this.
+	b := w.Bytes()
+	b = bytes.ReplaceAll(b, []byte(cwd), []byte("."))
+	cmd.OutOrStderr().Write(b)
+	if fatal {
+		exit()
 	}
 }
 

cmd/cue/cmd/root.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	logger "log"
 	"os"
 
@@ -43,7 +44,7 @@ var log = logger.New(os.Stderr, "", logger.Lshortfile)
 var cfgFile string
 
 // newRootCmd creates the base command when called without any subcommands
-func newRootCmd() *cobra.Command {
+func newRootCmd() *Command {
 	cmd := &cobra.Command{
 		Use:   "cue",
 		Short: "cue emits configuration files to user-defined commands.",
@@ -70,14 +71,16 @@ For more information on writing CUE configuration files see cuelang.org.`,
 		SilenceUsage: true,
 	}
 
+	cmdCmd := newCmdCmd()
+
 	subCommands := []*cobra.Command{
 		newTrimCmd(),
 		newImportCmd(),
 		newEvalCmd(),
 		newGetCmd(),
 		newFmtCmd(),
 		newExportCmd(),
-		newCmdCmd(),
+		cmdCmd,
 		newVetCmd(),
 		newAddCmd(),
 	}
@@ -88,28 +91,57 @@ For more information on writing CUE configuration files see cuelang.org.`,
 		cmd.AddCommand(sub)
 	}
 
-	return cmd
+	return &Command{root: cmd, cmd: cmdCmd}
 }
 
 // Main runs the cue tool. It loads the tool flags.
 func Main(ctx context.Context, args []string) (err error) {
+	cmd, err := New(args)
+	if err != nil {
+		return err
+	}
+	return cmd.Run(ctx)
+}
+
+type Command struct {
+	root *cobra.Command
+
+	// Subcommands
+	cmd *cobra.Command
+}
+
+func (c *Command) SetOutput(w io.Writer) {
+	c.root.SetOutput(w)
+}
+
+func (c *Command) Run(ctx context.Context) (err error) {
 	log.SetFlags(0)
 	// Three categories of commands:
 	// - normal
 	// - user defined
 	// - help
 	// For the latter two, we need to use the default loading.
-	defer func() {
-		switch e := recover().(type) {
-		case nil:
-		case panicError:
-			err = e.Err
-		default:
-			panic(err)
-		}
-		// We use panic to escape, instead of os.Exit
-	}()
-	rootCmd := newRootCmd()
+	defer recoverError(&err)
+
+	return c.root.Execute()
+}
+
+func recoverError(err *error) {
+	switch e := recover().(type) {
+	case nil:
+	case panicError:
+		*err = e.Err
+	default:
+		panic(e)
+	}
+	// We use panic to escape, instead of os.Exit
+}
+
+func New(args []string) (cmd *Command, err error) {
+	defer recoverError(&err)
+
+	cmd = newRootCmd()
+	rootCmd := cmd.root
 	rootCmd.SetArgs(args)
 	if len(args) >= 1 && args[0] != "help" {
 		// TODO: for now we only allow one instance. Eventually, we can allow
@@ -125,7 +157,7 @@ func Main(ctx context.Context, args []string) (err error) {
 			cmd  *cobra.Command
 		}
 		sub := map[string]subSpec{
-			"cmd": {commandSection, newCmdCmd()},
+			"cmd": {commandSection, cmd.cmd},
 			// "serve": {"server", nil},
 			// "fix":   {"fix", nil},
 		}
@@ -137,12 +169,12 @@ func Main(ctx context.Context, args []string) (err error) {
 				commands := tools.Lookup(sub.name)
 				i, err := commands.Fields()
 				if err != nil {
-					return err
+					return nil, err
 				}
 				for i.Next() {
 					_, _ = addCustom(sub.cmd, sub.name, i.Label(), tools)
 				}
-				return nil
+				return cmd, nil
 			}
 			tools := buildTools(rootCmd, args[1:])
 			_, err := addCustom(sub.cmd, sub.name, args[0], tools)
@@ -152,7 +184,7 @@ func Main(ctx context.Context, args []string) (err error) {
 			}
 		}
 	}
-	return rootCmd.Execute()
+	return cmd, nil
 }
 
 type panicError struct {

cmd/cue/cmd/vet.go

@@ -15,6 +15,8 @@
 package cmd
 
 import (
+	"os"
+
 	"cuelang.org/go/cue"
 	"cuelang.org/go/cue/ast"
 	"cuelang.org/go/cue/parser"
@@ -44,6 +46,8 @@ func doVet(cmd *cobra.Command, args []string) error {
 
 	w := cmd.OutOrStdout()
 
+	cwd, _ := os.Getwd()
+
 	for _, inst := range instances {
 		// TODO: use ImportPath or some other sanitized path.
 		opt := []cue.Option{
@@ -54,7 +58,7 @@ func doVet(cmd *cobra.Command, args []string) error {
 		}
 		err := inst.Value().Validate(opt...)
 		if flagVerbose.Bool(cmd) || err != nil {
-			printHeader(w, inst.Dir)
+			printHeader(w, cwd, inst.Dir)
 		}
 		exitIfErr(cmd, inst, err, false)
 	}

cmd/cue/main.go

@@ -16,7 +16,7 @@ package main
 
 import (
 	"context"
-	"log"
+	"fmt"
 	"os"
 
 	"cuelang.org/go/cmd/cue/cmd"
@@ -25,7 +25,7 @@ import (
 func main() {
 	err := cmd.Main(context.Background(), os.Args[1:])
 	if err != nil {
-		log.Println(err)
+		fmt.Println(err)
 		os.Exit(1)
 	}
 }