main_helper.go

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
}