Enable turning off the file header

arachnys · Mar 6, 2018 · 6df5af4 · 6df5af4
1 parent e7468fd
commit 6df5af4
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/cmd/protostub/cmd/generate.go b/cmd/protostub/cmd/generate.go
@@ -66,7 +66,7 @@ var generateCmd = &cobra.Command{
 			panic(err)
 		}
 
-		if err := gen.Gen(mf, p); err != nil {
+		if err := gen.Gen(mf, p, true); err != nil {
 			panic(err)
 		}
 	},

diff --git a/gen/gen.go b/gen/gen.go
@@ -21,12 +21,14 @@ type generator struct {
 // used to generate mypy stubs from proto data
 
 // Gen writes a mypy type stub to w, generated from the data in p
-func Gen(w io.Writer, p *protostub.ProtoData) error {
+func Gen(w io.Writer, p *protostub.ProtoData, header bool) error {
 	gen := &generator{0, bufio.NewWriter(w)}
 
-	gen.bw.WriteString(fmt.Sprintf("# Generated by protostub on %s\n", time.Now().Format("Mon Jan 2 15:04:05 -0700 MST 2006")))
-	gen.bw.WriteString("# https://github.com/arachnys/protostub\n")
-	gen.bw.WriteString("# Do not modify\n")
+	if header {
+		gen.bw.WriteString(fmt.Sprintf("# Generated by protostub on %s\n", time.Now().Format("Mon Jan 2 15:04:05 -0700 MST 2006")))
+		gen.bw.WriteString("# https://github.com/arachnys/protostub\n")
+		gen.bw.WriteString("# Do not modify\n")
+	}
 
 	for _, i := range p.Types {
 		if err := gen.gen(i); err != nil {

diff --git a/tests/gen_test.go b/tests/gen_test.go
@@ -54,7 +54,7 @@ func TestGeneration(t *testing.T) {
 
 		buf := bytes.NewBuffer(nil)
 
-		if err := gen.Gen(buf, p); err != nil {
+		if err := gen.Gen(buf, p, false); err != nil {
 			t.Fatal(err)
 		}