From 1cf8a7596127a0b196d8e40aa7fea569b73e1536 Mon Sep 17 00:00:00 2001
From: Tom Payne <twpayne@gmail.com>
Date: Thu, 19 Sep 2024 01:56:52 +0200
Subject: [PATCH] feat: add --file-header option

---
 cmd/gojsonstruct/main.go |  2 ++
 generator.go             | 11 +++++++++++
 generator_test.go        | 17 +++++++++++++++++
 3 files changed, 30 insertions(+)

diff --git a/cmd/gojsonstruct/main.go b/cmd/gojsonstruct/main.go
index 07670ab..905fe49 100644
--- a/cmd/gojsonstruct/main.go
+++ b/cmd/gojsonstruct/main.go
@@ -16,6 +16,7 @@ var (
 	abbreviations            = pflag.String("abbreviations", "", "comma-separated list of extra abbreviations")
 	format                   = pflag.String("format", "json", "format (json or yaml)")
 	decompress               = pflag.Bool("z", false, "decompress input with gzip")
+	fileHeader               = pflag.String("file-header", "", "file header")
 	omitEmptyTags            = pflag.String("omitempty-tags", "auto", "generate ,omitempty tags (never, always, or auto)")
 	packageComment           = pflag.String("package-comment", "", "package comment")
 	packageName              = pflag.String("package-name", "main", "package name")
@@ -40,6 +41,7 @@ func run() error {
 	pflag.Parse()
 
 	options := []jsonstruct.GeneratorOption{
+		jsonstruct.WithFileHeader(*fileHeader),
 		jsonstruct.WithOmitEmptyTags(omitEmptyTagsType[*omitEmptyTags]),
 		jsonstruct.WithSkipUnparsableProperties(*skipUnparsableProperties),
 		jsonstruct.WithStringTags(*stringTags),
diff --git a/generator.go b/generator.go
index 40a7564..7238208 100644
--- a/generator.go
+++ b/generator.go
@@ -32,6 +32,7 @@ type Generator struct {
 	abbreviations            map[string]bool
 	exportNameFunc           ExportNameFunc
 	exportRenames            map[string]string
+	fileHeader               string
 	goFormat                 bool
 	imports                  map[string]struct{}
 	intType                  string
@@ -76,6 +77,13 @@ func WithExtraAbbreviations(abbreviations ...string) GeneratorOption {
 	}
 }
 
+// WithFileHeader sets the file header.
+func WithFileHeader(fileHeader string) GeneratorOption {
+	return func(g *Generator) {
+		g.fileHeader = fileHeader
+	}
+}
+
 // WithGoFormat sets whether the output is should be formatted with go fmt.
 func WithGoFormat(goFormat bool) GeneratorOption {
 	return func(g *Generator) {
@@ -219,6 +227,9 @@ func NewGenerator(options ...GeneratorOption) *Generator {
 func (g *Generator) Generate() ([]byte, error) {
 	buffer := &bytes.Buffer{}
 	buffer.Grow(65536)
+	if g.fileHeader != "" {
+		fmt.Fprintf(buffer, "%s\n\n", g.fileHeader)
+	}
 	if g.packageComment != "" {
 		fmt.Fprintf(buffer, "// %s\n", g.packageComment)
 	}
diff --git a/generator_test.go b/generator_test.go
index 128e47a..d4f8568 100644
--- a/generator_test.go
+++ b/generator_test.go
@@ -998,6 +998,20 @@ func TestObserveJSONGoCode(t *testing.T) {
 				"// MyType is my type.\n" +
 				"type MyType bool\n",
 		},
+		{
+			name: "file_header",
+			json: "" +
+				`true`,
+			generatorOptions: []GeneratorOption{
+				WithFileHeader("// File header."),
+			},
+			expectedGoCodeStr: "" +
+				"// File header.\n" +
+				"\n" +
+				"package main\n" +
+				"\n" +
+				"type T bool\n",
+		},
 		{
 			name: "time",
 			json: `"1985-04-12T23:20:50.52Z"`,
@@ -1420,6 +1434,7 @@ func ExampleGenerator_ObserveJSONFile() {
 
 func ExampleGenerator_ObserveYAMLFile() {
 	generator := NewGenerator(
+		WithFileHeader("// Code generated by gojsonstruct. DO NOT EDIT."),
 		WithPackageName("mypackage"),
 		WithTypeName("MyType"),
 	)
@@ -1433,6 +1448,8 @@ func ExampleGenerator_ObserveYAMLFile() {
 	os.Stdout.Write(data)
 
 	// Output:
+	// // Code generated by gojsonstruct. DO NOT EDIT.
+	//
 	// package mypackage
 	//
 	// type MyType struct {