Format generated source code before writing it to the destination.
Signed-off-by: David Symonds <dsymonds@golang.org>
diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go
index 6371f7d..6698b2a 100644
--- a/mockgen/mockgen.go
+++ b/mockgen/mockgen.go
@@ -19,8 +19,10 @@
// 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 @@
packageName = "mock_" + sanitize(pkg.Name)
}
- g := generator{
- w: dst,
- }
+ g := new(generator)
if *source != "" {
g.filename = *source
} else {
@@ -98,6 +98,9 @@
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 @@
`
type generator struct {
- w io.Writer
+ buf bytes.Buffer
indent string
filename string // may be empty
@@ -133,7 +136,7 @@
}
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 @@
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
+}
