cmd/cue: add --outfile flag to get go

This introduces a way to write generated cue schemas to a given file or
stdout. It joins all declarations into one giant file and preserves all
package comments.

Author: fionera

cmd/cue/cmd/get_go.go

@@ -37,6 +37,7 @@ import (
 
 	cueast "cuelang.org/go/cue/ast"
 	"cuelang.org/go/cue/ast/astutil"
+	"cuelang.org/go/cue/errors"
 	"cuelang.org/go/cue/format"
 	"cuelang.org/go/cue/literal"
 	"cuelang.org/go/cue/load"
@@ -218,6 +219,8 @@ restrictive enum interpretation of #Switch remains.
 
 	cmd.Flags().StringP(string(flagPackage), "p", "", "package name for generated CUE files")
 
+	cmd.Flags().String(string(flagOutFile), "", "write generated CUE schema to the given file or stdout (only allows for one package)")
+
 	return cmd
 }
 
@@ -247,8 +250,9 @@ func (e *extractor) filter(name string) bool {
 type extractor struct {
 	cmd *Command
 
-	allPkgs map[string]*packages.Package
-	done    map[string]bool
+	allPkgs  map[string]*packages.Package
+	rootPkgs map[string]bool
+	done     map[string]bool
 
 	// per package
 	pkg      *packages.Package
@@ -357,6 +361,10 @@ var toString = []*types.Interface{
 // - consider not including types with any dropped fields.
 
 func extract(cmd *Command, args []string) error {
+	if flagLocal.IsSet(cmd) && flagOutFile.IsSet(cmd) {
+		return errors.New("--local and --outfile are mutually exclusive")
+	}
+
 	// TODO the CUE load using "." (below) assumes that a CUE module and a Go
 	// module will exist within the same directory (more precisely a Go module
 	// could be nested within a CUE module), such that the module path in any
@@ -387,9 +395,10 @@ func extract(cmd *Command, args []string) error {
 	}
 
 	e := extractor{
-		cmd:     cmd,
-		allPkgs: map[string]*packages.Package{},
-		orig:    map[types.Type]*ast.StructType{},
+		cmd:      cmd,
+		allPkgs:  map[string]*packages.Package{},
+		rootPkgs: map[string]bool{},
+		orig:     map[types.Type]*ast.StructType{},
 	}
 
 	e.initExclusions(flagExclude.String(cmd))
@@ -398,9 +407,14 @@ func extract(cmd *Command, args []string) error {
 
 	for _, p := range pkgs {
 		e.done[p.PkgPath] = true
+		e.rootPkgs[p.PkgPath] = true
 		e.addPackage(p)
 	}
 
+	if flagOutFile.IsSet(cmd) && len(e.rootPkgs) > 1 {
+		return errors.New("--outfile only allows for one package to be specified")
+	}
+
 	for _, p := range pkgs {
 		if err := e.extractPkg(root, p); err != nil {
 			return err
@@ -439,6 +453,10 @@ func (e *extractor) extractPkg(root string, p *packages.Package) error {
 
 	e.recordTypeInfo(p)
 
+	if !e.rootPkgs[p.PkgPath] && flagOutFile.IsSet(e.cmd) {
+		return nil
+	}
+
 	e.consts = map[string][]string{}
 
 	for _, f := range p.Syntax {
@@ -462,8 +480,10 @@ func (e *extractor) extractPkg(root string, p *packages.Package) error {
 		}
 	}
 
-	if err := os.MkdirAll(dir, 0777); err != nil {
-		return err
+	if !flagOutFile.IsSet(e.cmd) {
+		if err := os.MkdirAll(dir, 0777); err != nil {
+			return err
+		}
 	}
 
 	e.usedPkgs = map[string]bool{}
@@ -473,7 +493,21 @@ func (e *extractor) extractPkg(root string, p *packages.Package) error {
 		args += " --exclude=" + e.exclude
 	}
 
+	pName := flagPackage.String(e.cmd)
+	if pName == "" {
+		pName = p.Name
+	}
+
+	var (
+		decls  []cueast.Decl
+		cuePkg *cueast.Package
+	)
 	for i, f := range p.Syntax {
+		if cuePkg == nil || !flagOutFile.IsSet(e.cmd) {
+			cuePkg = &cueast.Package{Name: e.ident(pName, false)}
+			decls = nil
+		}
+
 		e.cmap = ast.NewCommentMap(p.Fset, f, f.Comments)
 
 		e.pkgNames = map[string]pkgInfo{}
@@ -494,25 +528,24 @@ func (e *extractor) extractPkg(root string, p *packages.Package) error {
 			e.pkgNames[pkgPath] = info
 		}
 
-		decls := []cueast.Decl{}
+		var fileDecls []cueast.Decl
 		for _, d := range f.Decls {
 			switch d := d.(type) {
 			case *ast.GenDecl:
-				decls = append(decls, e.reportDecl(d)...)
+				fileDecls = append(fileDecls, e.reportDecl(d)...)
 			}
 		}
 
-		if len(decls) == 0 && f.Doc == nil {
+		if len(fileDecls) == 0 && f.Doc == nil && !flagOutFile.IsSet(e.cmd) {
 			continue
 		}
+		decls = append(decls, fileDecls...)
 
-		pName := flagPackage.String(e.cmd)
-		if pName == "" {
-			pName = p.Name
-		}
+		addDoc(f.Doc, cuePkg)
 
-		pkg := &cueast.Package{Name: e.ident(pName, false)}
-		addDoc(f.Doc, pkg)
+		if flagOutFile.IsSet(e.cmd) && i != len(p.Syntax)-1 {
+			continue
+		}
 
 		f := &cueast.File{Decls: []cueast.Decl{
 			&cueast.CommentGroup{List: []*cueast.Comment{
@@ -521,7 +554,7 @@ func (e *extractor) extractPkg(root string, p *packages.Package) error {
 			&cueast.CommentGroup{List: []*cueast.Comment{
 				{Text: "//cue:generate cue get go " + args},
 			}},
-			pkg,
+			cuePkg,
 		}}
 		f.Decls = append(f.Decls, decls...)
 
@@ -530,16 +563,29 @@ func (e *extractor) extractPkg(root string, p *packages.Package) error {
 		}
 
 		file := filepath.Base(p.CompiledGoFiles[i])
-
 		file = strings.Replace(file, ".go", "_go", 1)
 		file += "_gen.cue"
+		filePath := filepath.Join(dir, file)
+
 		b, err := format.Node(f, format.Simplify())
 		if err != nil {
 			return err
 		}
-		err = os.WriteFile(filepath.Join(dir, file), b, 0666)
-		if err != nil {
-			return err
+
+		if flagOutFile.IsSet(e.cmd) {
+			filePath = flagOutFile.String(e.cmd)
+		}
+
+		if filePath == "-" {
+			_, err = os.Stdout.Write(b)
+			if err != nil {
+				return err
+			}
+		} else {
+			err = os.WriteFile(filePath, b, 0666)
+			if err != nil {
+				return err
+			}
 		}
 	}
 
@@ -549,10 +595,10 @@ func (e *extractor) extractPkg(root string, p *packages.Package) error {
 		}
 	}
 
-	for path := range e.usedPkgs {
-		if !e.done[path] {
-			e.done[path] = true
-			if err := e.extractPkg(root, e.allPkgs[path]); err != nil {
+	for pkgPath := range e.usedPkgs {
+		if !e.done[pkgPath] {
+			e.done[pkgPath] = true
+			if err := e.extractPkg(root, e.allPkgs[pkgPath]); err != nil {
 				return err
 			}
 		}
@@ -883,7 +929,7 @@ func makeDoc(g *ast.CommentGroup, isDoc bool) *cueast.CommentGroup {
 		// The parser has given us exactly the comment text.
 		switch c[1] {
 		case '/':
-			//-style comment (no newline at the end)
+			// -style comment (no newline at the end)
 			a = append(a, &cueast.Comment{Text: c})
 
 		case '*':