diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go index 6371f7da..6698b2a8 100644 --- a/mockgen/mockgen.go +++ b/mockgen/mockgen.go @@ -19,8 +19,10 @@ package main // TODO: This does not support embedding package-local interfaces in a separate file. import ( + "bytes" "flag" "fmt" + "go/format" "go/token" "io" "log" @@ -86,9 +88,7 @@ func main() { packageName = "mock_" + sanitize(pkg.Name) } - g := generator{ - w: dst, - } + g := new(generator) if *source != "" { g.filename = *source } else { @@ -98,6 +98,9 @@ func main() { if err := g.Generate(pkg, packageName); err != nil { log.Fatalf("Failed generating mock: %v", err) } + if _, err := dst.Write(g.Output()); err != nil { + log.Fatalf("Failed writing to destination: %v", err) + } } func usage() { @@ -123,7 +126,7 @@ Example: ` type generator struct { - w io.Writer + buf bytes.Buffer indent string filename string // may be empty @@ -133,7 +136,7 @@ type generator struct { } func (g *generator) p(format string, args ...interface{}) { - fmt.Fprintf(g.w, g.indent+format+"\n", args...) + fmt.Fprintf(&g.buf, g.indent+format+"\n", args...) } func (g *generator) in() { @@ -411,3 +414,12 @@ func (g *generator) GenerateMockRecorderMethod(mockType string, m *model.Method) g.p("}") return nil } + +// Output returns the generator's output, formatted in the standard Go style. +func (g *generator) Output() []byte { + src, err := format.Source(g.buf.Bytes()) + if err != nil { + log.Fatalf("Failed to format generated source code: %s\n%s", err, g.buf.String()) + } + return src +}