Simple go code generator for fitasks

Based on the canonical example in the go code

upup/tools/generators/fitask/generator.go

@@ -0,0 +1,84 @@
+package main
+
+import (
+	"io"
+	"k8s.io/kube-deploy/upup/tools/generators/pkg/codegen"
+	"text/template"
+)
+
+type FitaskGenerator struct {
+	Package string
+}
+
+var _ codegen.Generator = &FitaskGenerator{}
+
+const headerDef = `
+package {{.Package}}
+
+import (
+	"encoding/json"
+
+	"k8s.io/kube-deploy/upup/pkg/fi"
+)
+`
+
+const perTypeDef = `
+// {{.Name}}
+
+// JSON marshalling boilerplate
+type real{{.Name}} {{.Name}}
+
+func (o *{{.Name}}) UnmarshalJSON(data []byte) error {
+	var jsonName string
+	if err := json.Unmarshal(data, &jsonName); err == nil {
+		o.Name = &jsonName
+		return nil
+	}
+
+	var r real{{.Name}}
+	if err := json.Unmarshal(data, &r); err != nil {
+		return err
+	}
+	*o = {{.Name}}(r)
+	return nil
+}
+
+var _ fi.HasName = &{{.Name}}{}
+
+func (e *{{.Name}}) GetName() *string {
+	return e.Name
+}
+
+func (e *{{.Name}}) SetName(name string) {
+	e.Name = &name
+}
+
+func (e *{{.Name}}) String() string {
+	return fi.TaskAsString(e)
+}
+`
+
+func (g *FitaskGenerator) Init(parser *codegen.GoParser) error {
+	g.Package = parser.Package.Name
+
+	return nil
+}
+
+func (g *FitaskGenerator) WriteHeader(w io.Writer) error {
+	t := template.Must(template.New("Header").Parse(headerDef))
+
+	return t.Execute(w, g)
+}
+
+type TypeData struct {
+	Name string
+}
+
+func (g *FitaskGenerator) WriteType(w io.Writer, typeName string) error {
+	t := template.Must(template.New("PerType").Parse(perTypeDef))
+
+	d := &TypeData{}
+	d.Name = typeName
+
+	return t.Execute(w, d)
+}

upup/tools/generators/fitask/main.go

@@ -0,0 +1,10 @@
+package main
+
+import (
+	"k8s.io/kube-deploy/upup/tools/generators/pkg/codegen"
+)
+
+func main() {
+	generator := &FitaskGenerator{}
+	codegen.RunGenerator("fitask", generator)
+}

upup/tools/generators/pkg/codegen/generator.go

@@ -0,0 +1,9 @@
+package codegen
+
+import "io"
+
+type Generator interface {
+	Init(parser *GoParser) error
+	WriteHeader(w io.Writer) error
+	WriteType(w io.Writer, typeName string) error
+}

upup/tools/generators/pkg/codegen/main_helper.go

@@ -0,0 +1,115 @@
+package codegen
+
+import (
+	"bytes"
+	"flag"
+	"fmt"
+	"go/format"
+	"io/ioutil"
+	"log"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+var (
+	flagTypes = flag.String("type", "", "comma-separated list of type names; must be set")
+	output    = flag.String("output", "", "output file name; default srcdir/<type>_fitask.go")
+)
+
+// Usage is a replacement usage function for the flags package.
+func Usage() {
+	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
+	fmt.Fprintf(os.Stderr, "\tfitask [flags] -type T [directory]\n")
+	fmt.Fprintf(os.Stderr, "\tfitask [flags] -type T files... # Must be a single package\n")
+	fmt.Fprintf(os.Stderr, "Flags:\n")
+	flag.PrintDefaults()
+}
+
+// isDirectory reports whether the named file is a directory.
+func isDirectory(name string) bool {
+	info, err := os.Stat(name)
+	if err != nil {
+		log.Fatal(err)
+	}
+	return info.IsDir()
+}
+
+func RunGenerator(key string, generator Generator) {
+	flag.Usage = Usage
+	flag.Parse()
+	if len(*flagTypes) == 0 {
+		flag.Usage()
+		os.Exit(2)
+	}
+	typeNames := strings.Split(*flagTypes, ",")
+
+	// We accept either one directory or a list of files. Which do we have?
+	args := flag.Args()
+	if len(args) == 0 {
+		// Default: process whole package in current directory.
+		args = []string{"."}
+	}
+
+	// Parse the package once.
+	var dir string
+	var parser GoParser
+
+	if len(args) == 1 && isDirectory(args[0]) {
+		dir = args[0]
+		parser.parsePackageDir(args[0])
+		//} else {
+		//	dir = filepath.Dir(args[0])
+		//	g.parsePackageFiles(args)
+	}
+
+	err := generator.Init(&parser)
+	if err != nil {
+		log.Fatalf("error initializing generator: %v", err)
+	}
+
+	var b bytes.Buffer
+
+	fmt.Fprintf(&b, "// Code generated by \"%q %s\"; DO NOT EDIT\n\n", key, strings.Join(os.Args[1:], " "))
+
+	err = generator.WriteHeader(&b)
+	if err != nil {
+		log.Fatalf("error generating header: %v", err)
+	}
+
+	// Run generate for each type.
+	for _, typeName := range typeNames {
+		err := generator.WriteType(&b, typeName)
+		if err != nil {
+			log.Fatalf("error generating header: %v", err)
+		}
+	}
+
+	// gofmt the output.
+	src := gofmt(b.Bytes())
+
+	// Write to file.
+	outputName := *output
+	if outputName == "" {
+		baseName := fmt.Sprintf("%s_%s.go", typeNames[0], key)
+		outputName = filepath.Join(dir, strings.ToLower(baseName))
+	}
+	err = ioutil.WriteFile(outputName, src, 0644)
+	if err != nil {
+		log.Fatalf("writing output: %s", err)
+	}
+
+}
+
+// format returns the gofmt-ed contents of the Generator's buffer.
+func gofmt(src []byte) []byte {
+	formatted, err := format.Source(src)
+	if err != nil {
+		// Should never happen, but can arise when developing this code.
+		// The user can compile the output to see the error.
+		log.Printf("warning: internal error: invalid Go generated: %s", err)
+		log.Printf("warning: compile the package to analyze the error")
+		return src
+	}
+	return formatted
+}

upup/tools/generators/pkg/codegen/parse.go

@@ -0,0 +1,106 @@
+package codegen
+
+import (
+	"go/ast"
+	"go/build"
+	"go/importer"
+	"go/parser"
+	"go/token"
+	"go/types"
+	"log"
+	"path/filepath"
+	"strings"
+)
+
+type GoParser struct {
+	Package *Package // Package we are scanning.
+}
+
+// File holds a single parsed file and associated data.
+type File struct {
+	pkg  *Package  // Package to which this file belongs.
+	file *ast.File // Parsed AST.
+	// These fields are reset for each type being generated.
+	typeName string // Name of the constant type.
+	//values   []Value // Accumulator for constant values of that type.
+}
+
+type Package struct {
+	dir      string
+	Name     string
+	defs     map[*ast.Ident]types.Object
+	files    []*File
+	typesPkg *types.Package
+}
+
+// prefixDirectory prepends the directory name on the beginning of each name in the list.
+func prefixDirectory(directory string, names []string) []string {
+	if directory == "." {
+		return names
+	}
+	ret := make([]string, len(names))
+	for i, name := range names {
+		ret[i] = filepath.Join(directory, name)
+	}
+	return ret
+}
+
+// parsePackageDir parses the package residing in the directory.
+func (g *GoParser) parsePackageDir(directory string) {
+	pkg, err := build.Default.ImportDir(directory, 0)
+	if err != nil {
+		log.Fatalf("cannot process directory %s: %s", directory, err)
+	}
+	var names []string
+	names = append(names, pkg.GoFiles...)
+	//names = append(names, pkg.CgoFiles...)
+	//names = append(names, pkg.SFiles...)
+	names = prefixDirectory(directory, names)
+	g.parsePackage(directory, names, nil)
+}
+
+// parsePackage analyzes the single package constructed from the named files.
+// If text is non-nil, it is a string to be used instead of the content of the file,
+// to be used for testing. parsePackage exits if there is an error.
+func (g *GoParser) parsePackage(directory string, names []string, text interface{}) {
+	var files []*File
+	var astFiles []*ast.File
+	g.Package = new(Package)
+	fs := token.NewFileSet()
+	for _, name := range names {
+		if !strings.HasSuffix(name, ".go") {
+			continue
+		}
+		parsedFile, err := parser.ParseFile(fs, name, text, 0)
+		if err != nil {
+			log.Fatalf("parsing package: %s: %s", name, err)
+		}
+		astFiles = append(astFiles, parsedFile)
+		files = append(files, &File{
+			file: parsedFile,
+			pkg:  g.Package,
+		})
+	}
+	if len(astFiles) == 0 {
+		log.Fatalf("%s: no buildable Go files", directory)
+	}
+	g.Package.Name = astFiles[0].Name.Name
+	g.Package.files = files
+	g.Package.dir = directory
+	// Type check the package.
+	g.Package.check(fs, astFiles)
+}
+
+// check type-checks the package. The package must be OK to proceed.
+func (pkg *Package) check(fs *token.FileSet, astFiles []*ast.File) {
+	pkg.defs = make(map[*ast.Ident]types.Object)
+	config := types.Config{Importer: importer.Default(), FakeImportC: true}
+	info := &types.Info{
+		Defs: pkg.defs,
+	}
+	typesPkg, err := config.Check(pkg.dir, fs, astFiles, info)
+	if err != nil {
+		log.Fatalf("checking package: %s", err)
+	}
+	pkg.typesPkg = typesPkg
+}