Enable turning off the file header

Author: Will Huxtable

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)
 		}
 	},

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 {

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)
 		}