diff --git a/cmd/protostub/cmd/generate.go b/cmd/protostub/cmd/generate.go index 30ebfd3..74aca9a 100644 --- 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 index cbace08..84731ca 100644 --- 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 index c08c49d..db97ea0 100644 --- 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) }