// Flattens a main file with rewritten dependencies.
func flatten(pkg string, main []byte, deps [][]byte, remDead bool) ([]byte, error) {
buffer := new(bytes.Buffer)
// Dump all code pieces into a buffer
fmt.Fprintf(buffer, "package %s\n\n", pkg)
fmt.Fprintf(buffer, "%s\n\n", main)
for _, dep := range deps {
fmt.Fprintf(buffer, "%s\n\n", dep)
}
// Format the blob to Go standards and add imports
blob, err := imports.Process("", buffer.Bytes(), nil)
if err != nil {
return nil, err
}
// Shake all dependencies, format and goimport again
if remDead {
if blob, err = shake(blob); err != nil {
return nil, err
}
if blob, err = imports.Process("", blob, nil); err != nil {
return nil, err
}
}
return blob, nil
}
func main() {
file := os.Args[1]
if !strings.HasSuffix(file, ".go") || strings.HasSuffix(file, "_test.go") {
log.Fatal("arg must be go file, not test")
}
in, err := Parse(file)
if err != nil {
log.Fatal(err)
}
var buf bytes.Buffer
tmpl.Execute(&buf, in)
src, err := imports.Process("", buf.Bytes(), nil)
if err != nil {
log.Fatalf("%s\nfailed to gofmt generated code: %v", src, err)
}
file = strings.TrimSuffix(file, ".go") + "_test.go"
out, err := os.OpenFile(file, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0644)
if err != nil {
log.Fatal(err)
}
defer out.Close()
if _, err := out.Write(src); err != nil {
log.Fatal(err)
}
}
func goimports(src []byte) (string, string) {
out, err := imports.Process("", src, nil)
if err != nil {
return "", fmt.Sprintf("%v", err)
}
return string(out), ""
}
// findInterface returns the import path and identifier of an interface.
// For example, given "http.ResponseWriter", findInterface returns
// "net/http", "ResponseWriter".
func findInterface(iface string) (path string, id string, err error) {
if len(strings.Fields(iface)) != 1 {
return "", "", fmt.Errorf("couldn't parse interface: %s", iface)
}
// Let goimports do the heavy lifting.
src := []byte("package hack\n" + "var i " + iface)
imp, err := imports.Process("", src, nil)
if err != nil {
return "", "", fmt.Errorf("couldn't parse interface: %s", iface)
}
// imp should now contain an appropriate import.
// Parse out the import and the identifier.
fset := token.NewFileSet()
f, err := parser.ParseFile(fset, "", imp, 0)
if err != nil {
panic(err)
}
if len(f.Imports) == 0 {
return "", "", fmt.Errorf("unrecognized interface: %s", iface)
}
raw := f.Imports[0].Path.Value // "io"
path, err = strconv.Unquote(raw) // io
if err != nil {
panic(err)
}
decl := f.Decls[1].(*ast.GenDecl) // var i io.Reader
spec := decl.Specs[0].(*ast.ValueSpec) // i io.Reader
sel := spec.Type.(*ast.SelectorExpr) // io.Reader
id = sel.Sel.Name // Reader
return path, id, nil
}
func main() {
iface := "http.Handler"
src := "package hack; var i " + iface // HL
fmt.Println(src, "\n---")
imp, _ := imports.Process("", []byte(src), nil) // HL
// ignoring errors throughout this presentation
fmt.Println(string(imp))
}
// WriteAll writes the generated code for all Types and TypeWriters in the App to respective files.
func (a *app) WriteAll() error {
// TypeWriters which will write for each type; use string as key because Type is not comparable
var writersByType = make(map[string][]TypeWriter)
// validate them all (don't fail halfway)
for _, t := range a.Types {
for _, tw := range a.TypeWriters {
write, err := tw.Validate(t)
if err != nil {
return err
}
if write {
writersByType[t.String()] = append(writersByType[t.String()], tw)
}
}
}
// one buffer for each file, keyed by file name
buffers := make(map[string]*bytes.Buffer)
// write the generated code for each Type & TypeWriter into memory
for _, t := range a.Types {
for _, tw := range writersByType[t.String()] {
var b bytes.Buffer
write(&b, a, t, tw)
// will append _test to file name if the source type is in a _test.go file
f := strings.ToLower(fmt.Sprintf("%s_%s%s.go", t.Name, tw.Name(), t.test))
buffers[f] = &b
}
}
// validate generated ast's before committing to files
for f, b := range buffers {
if _, err := parser.ParseFile(token.NewFileSet(), f, b.String(), 0); err != nil {
// TODO: prompt to write (ignored) _file on error? parsing errors are meaningless without.
return err
}
}
// format, remove unused imports, and commit to files
for f, b := range buffers {
src, err := imports.Process(f, b.Bytes(), nil)
// shouldn't be an error if the ast parsing above succeeded
if err != nil {
return err
}
if err := writeFile(f, src); err != nil {
return err
}
}
return nil
}
func processFile(filename string, in io.Reader, out io.Writer, stdin bool) error {
opt := options
if stdin {
nopt := *options
nopt.Fragment = true
opt = &nopt
}
if in == nil {
f, err := os.Open(filename)
if err != nil {
return err
}
defer f.Close()
in = f
}
src, err := ioutil.ReadAll(in)
if err != nil {
return err
}
res, err := imports.Process(filename, src, opt)
if err != nil {
return err
}
if !bytes.Equal(src, res) {
// formatting has changed
if *list {
fmt.Fprintln(out, filename)
}
if *write {
err = ioutil.WriteFile(filename, res, 0)
if err != nil {
return err
}
}
if *doDiff {
data, err := diff(src, res)
if err != nil {
return fmt.Errorf("computing diff: %s", err)
}
fmt.Printf("diff %s gofmt/%s\n", filename, filename)
out.Write(data)
}
}
if !*list && !*write && !*doDiff {
_, err = out.Write(res)
}
return err
}
func generateFile(h *HitType) {
code := &bytes.Buffer{}
codeTemplate.Execute(code, h)
newCode, err := imports.Process("prog.go", code.Bytes(), &imports.Options{
AllErrors: true,
TabWidth: 4,
Comments: true,
})
if err != nil {
displayErr(code.String(), err)
}
fname := "type-" + h.Name + ".go"
f, err := os.Create(fname)
defer f.Close()
if err != nil {
log.Fatal(err)
}
f.Write(newCode)
log.Println("generated " + fname)
}