Repository: abice/go-enum
Pull request #53 (2021-04-29 19:57:30)

Add --template flag to provide filenames of additional templates
The `--template`/`-t` flag allows users to provide their own templates and extend code generation.

https://github.com/abice/go-enum/issues/52


Patch:
diff --git a/README.md b/README.md
index 3bbde58..5f6b026 100644
--- a/README.md
+++ b/README.md
@@ -26,20 +26,21 @@ USAGE:
    go-enum [global options] [arguments...]
 
 GLOBAL OPTIONS:
-   --file value, -f value  The file(s) to generate enums.  Use more than one flag for more files.
-   --noprefix              Prevents the constants generated from having the Enum as a prefix. (default: false)
-   --lower                 Adds lowercase variants of the enum strings for lookup. (default: false)
-   --nocase                Adds case insensitive parsing to the enumeration (forces lower flag). (default: false)
-   --marshal               Adds text (and inherently json) marshalling functions. (default: false)
-   --sql                   Adds SQL database scan and value functions. (default: false)
-   --flag                  Adds golang flag functions. (default: false)
-   --prefix value          Replaces the prefix with a user one.
-   --names                 Generates a 'Names() []string' function, and adds the possible enum values in the error response during parsing (default: false)
-   --nocamel               Removes the snake_case to CamelCase name changing (default: false)
-   --ptr                   Adds a pointer method to get a pointer from const values (default: false)
-   --sqlnullint            Adds a Null{{ENUM}} type for marshalling a nullable int value to sql (default: false)
-   --sqlnullstr            Adds a Null{{ENUM}} type for marshalling a nullable string value to sql.  If sqlnullint is specified too, it will be Null{{ENUM}}Str (default: false)
-   --help, -h              show help (default: false)
+   --file value, -f value      The file(s) to generate enums.  Use more than one flag for more files.
+   --noprefix                  Prevents the constants generated from having the Enum as a prefix. (default: false)
+   --lower                     Adds lowercase variants of the enum strings for lookup. (default: false)
+   --nocase                    Adds case insensitive parsing to the enumeration (forces lower flag). (default: false)
+   --marshal                   Adds text (and inherently json) marshalling functions. (default: false)
+   --sql                       Adds SQL database scan and value functions. (default: false)
+   --flag                      Adds golang flag functions. (default: false)
+   --prefix value              Replaces the prefix with a user one.
+   --names                     Generates a 'Names() []string' function, and adds the possible enum values in the error response during parsing (default: false)
+   --nocamel                   Removes the snake_case to CamelCase name changing (default: false)
+   --ptr                       Adds a pointer method to get a pointer from const values (default: false)
+   --sqlnullint                Adds a Null{{ENUM}} type for marshalling a nullable int value to sql (default: false)
+   --sqlnullstr                Adds a Null{{ENUM}} type for marshalling a nullable string value to sql.  If sqlnullint is specified too, it will be Null{{ENUM}}Str (default: false)
+   --template value, -t value  Additional template file(s) to generate enums.  Use more than one flag for more files. Templates will be executed in alphabetical order.
+   --help, -h                  show help (default: false)
 ```
 
 ### Syntax
diff --git a/example/user_globbed.tmpl b/example/user_globbed.tmpl
new file mode 100644
index 0000000..a02e277
--- /dev/null
+++ b/example/user_globbed.tmpl
@@ -0,0 +1,3 @@
+func Parse{{.enum.Name}}GlobbedExample() bool {
+    return true
+}
diff --git a/example/user_globbed2.tmpl b/example/user_globbed2.tmpl
new file mode 100644
index 0000000..eae0a54
--- /dev/null
+++ b/example/user_globbed2.tmpl
@@ -0,0 +1,3 @@
+func Parse{{.enum.Name}}GlobbedExample2() bool {
+    return true
+}
diff --git a/example/user_template.go b/example/user_template.go
new file mode 100644
index 0000000..6527071
--- /dev/null
+++ b/example/user_template.go
@@ -0,0 +1,13 @@
+//go:generate ../bin/go-enum -f=$GOFILE -t user_template.tmpl -t *user_glob*.tmpl
+
+package example
+
+// OceanColor is an enumeration of ocean colors that are allowed.
+/*
+ENUM(
+Cerulean
+Blue
+Green
+)
+*/
+type OceanColor int
diff --git a/example/user_template.tmpl b/example/user_template.tmpl
new file mode 100644
index 0000000..7aab62b
--- /dev/null
+++ b/example/user_template.tmpl
@@ -0,0 +1,4 @@
+// Additional template
+func Parse{{.enum.Name}}Example() bool {
+    return true
+}
diff --git a/example/user_template_enum.go b/example/user_template_enum.go
new file mode 100644
index 0000000..f5d8ed9
--- /dev/null
+++ b/example/user_template_enum.go
@@ -0,0 +1,59 @@
+// Code generated by go-enum
+// DO NOT EDIT!
+
+package example
+
+import (
+	"fmt"
+)
+
+const (
+	// OceanColorCerulean is a OceanColor of type Cerulean.
+	OceanColorCerulean OceanColor = iota
+	// OceanColorBlue is a OceanColor of type Blue.
+	OceanColorBlue
+	// OceanColorGreen is a OceanColor of type Green.
+	OceanColorGreen
+)
+
+const _OceanColorName = "CeruleanBlueGreen"
+
+var _OceanColorMap = map[OceanColor]string{
+	0: _OceanColorName[0:8],
+	1: _OceanColorName[8:12],
+	2: _OceanColorName[12:17],
+}
+
+// String implements the Stringer interface.
+func (x OceanColor) String() string {
+	if str, ok := _OceanColorMap[x]; ok {
+		return str
+	}
+	return fmt.Sprintf("OceanColor(%d)", x)
+}
+
+var _OceanColorValue = map[string]OceanColor{
+	_OceanColorName[0:8]:   0,
+	_OceanColorName[8:12]:  1,
+	_OceanColorName[12:17]: 2,
+}
+
+// ParseOceanColor attempts to convert a string to a OceanColor
+func ParseOceanColor(name string) (OceanColor, error) {
+	if x, ok := _OceanColorValue[name]; ok {
+		return x, nil
+	}
+	return OceanColor(0), fmt.Errorf("%s is not a valid OceanColor", name)
+}
+
+func ParseOceanColorGlobbedExample() bool {
+	return true
+}
+func ParseOceanColorGlobbedExample2() bool {
+	return true
+}
+
+// Additional template
+func ParseOceanColorExample() bool {
+	return true
+}
diff --git a/generator/generator.go b/generator/generator.go
index a929c83..db81498 100644
--- a/generator/generator.go
+++ b/generator/generator.go
@@ -28,21 +28,22 @@ const (
 
 // Generator is responsible for generating validation files for the given in a go source file.
 type Generator struct {
-	t               *template.Template
-	knownTemplates  map[string]*template.Template
-	fileSet         *token.FileSet
-	noPrefix        bool
-	lowercaseLookup bool
-	caseInsensitive bool
-	marshal         bool
-	sql             bool
-	flag            bool
-	names           bool
-	leaveSnakeCase  bool
-	prefix          string
-	sqlNullInt      bool
-	sqlNullStr      bool
-	ptr             bool
+	t                 *template.Template
+	knownTemplates    map[string]*template.Template
+	userTemplateNames []string
+	fileSet           *token.FileSet
+	noPrefix          bool
+	lowercaseLookup   bool
+	caseInsensitive   bool
+	marshal           bool
+	sql               bool
+	flag              bool
+	names             bool
+	leaveSnakeCase    bool
+	prefix            string
+	sqlNullInt        bool
+	sqlNullStr        bool
+	ptr               bool
 }
 
 // Enum holds data for a discovered enum in the parsed source
@@ -66,10 +67,11 @@ type EnumValue struct {
 // templates loaded.
 func NewGenerator() *Generator {
 	g := &Generator{
-		knownTemplates: make(map[string]*template.Template),
-		t:              template.New("generator"),
-		fileSet:        token.NewFileSet(),
-		noPrefix:       false,
+		knownTemplates:    make(map[string]*template.Template),
+		userTemplateNames: make([]string, 0),
+		t:                 template.New("generator"),
+		fileSet:           token.NewFileSet(),
+		noPrefix:          false,
 	}
 
 	funcs := sprig.TxtFuncMap()
@@ -163,6 +165,18 @@ func (g *Generator) WithSQLNullStr() *Generator {
 	return g
 }
 
+// WithTemplates is used to provide the filenames of additional templates.
+func (g *Generator) WithTemplates(filenames ...string) *Generator {
+	for _, ut := range template.Must(g.t.ParseFiles(filenames...)).Templates() {
+		if _, ok := g.knownTemplates[ut.Name()]; !ok {
+			g.userTemplateNames = append(g.userTemplateNames, ut.Name())
+		}
+	}
+	g.updateTemplates()
+	sort.Strings(g.userTemplateNames)
+	return g
+}
+
 // GenerateFromFile is responsible for orchestrating the Code generation.  It results in a byte array
 // that can be written to any file desired.  It has already had goimports run on the code before being returned.
 func (g *Generator) GenerateFromFile(inputFile string) ([]byte, error) {
@@ -223,6 +237,13 @@ func (g *Generator) Generate(f *ast.File) ([]byte, error) {
 		if err != nil {
 			return vBuff.Bytes(), errors.WithMessage(err, fmt.Sprintf("Failed writing enum data for enum: %q", name))
 		}
+
+		for _, userTemplateName := range g.userTemplateNames {
+			err = g.t.ExecuteTemplate(vBuff, userTemplateName, data)
+			if err != nil {
+				return vBuff.Bytes(), errors.WithMessage(err, fmt.Sprintf("Failed writing enum data for enum: %q, template: %v", name, userTemplateName))
+			}
+		}
 	}
 
 	formatted, err := imports.Process(pkg, vBuff.Bytes(), nil)
diff --git a/main.go b/main.go
index ff1623e..4910227 100644
--- a/main.go
+++ b/main.go
@@ -14,19 +14,20 @@ import (
 )
 
 type rootT struct {
-	FileNames      cli.StringSlice
-	NoPrefix       bool
-	Lowercase      bool
-	NoCase         bool
-	Marshal        bool
-	SQL            bool
-	Flag           bool
-	Prefix         string
-	Names          bool
-	LeaveSnakeCase bool
-	SQLNullStr     bool
-	SQLNullInt     bool
-	Ptr            bool
+	FileNames         cli.StringSlice
+	NoPrefix          bool
+	Lowercase         bool
+	NoCase            bool
+	Marshal           bool
+	SQL               bool
+	Flag              bool
+	Prefix            string
+	Names             bool
+	LeaveSnakeCase    bool
+	SQLNullStr        bool
+	SQLNullInt        bool
+	Ptr               bool
+	TemplateFileNames cli.StringSlice
 }
 
 func main() {
@@ -109,6 +110,12 @@ func main() {
 				Usage:       "Adds a Null{{ENUM}} type for marshalling a nullable string value to sql.  If sqlnullint is specified too, it will be Null{{ENUM}}Str",
 				Destination: &argv.SQLNullStr,
 			},
+			&cli.StringSliceFlag{
+				Name:        "template",
+				Aliases:     []string{"t"},
+				Usage:       "Additional template file(s) to generate enums.  Use more than one flag for more files. Templates will be executed in alphabetical order.",
+				Destination: &argv.TemplateFileNames,
+			},
 		},
 		Action: func(ctx *cli.Context) error {
 			for _, fileOption := range argv.FileNames.Value() {
@@ -151,20 +158,21 @@ func main() {
 				if argv.SQLNullStr {
 					g.WithSQLNullStr()
 				}
+				if templates := []string(argv.TemplateFileNames.Value()); len(templates) > 0 {
+					for _, t := range templates {
+						if fn, err := globFilenames(t); err != nil {
+							return err
+						} else {
+							g.WithTemplates(fn...)
+						}
+					}
+				}
 
 				var filenames []string
-
-				// In order to maintain existing capabilities, only glob when a * is in the path.
-				// Leave execution on par with old method in case there are bad patterns in use that somehow
-				// work without the Glob method.
-				if strings.Contains(fileOption, "*") {
-					matches, err := filepath.Glob(fileOption)
-					if err != nil {
-						return fmt.Errorf("failed parsing glob filepath\nInputFile=%s\nError=%s", color.Cyan(fileOption), color.RedBg(err))
-					}
-					filenames = append(filenames, matches...)
+				if fn, err := globFilenames(fileOption); err != nil {
+					return err
 				} else {
-					filenames = append(filenames, fileOption)
+					filenames = fn
 				}
 
 				for _, fileName := range filenames {
@@ -203,3 +211,19 @@ func main() {
 		log.Fatal(err)
 	}
 }
+
+// globFilenames gets a list of filenames matching the provided filename.
+// In order to maintain existing capabilities, only glob when a * is in the path.
+// Leave execution on par with old method in case there are bad patterns in use that somehow
+// work without the Glob method.
+func globFilenames(filename string) ([]string, error) {
+	if strings.Contains(filename, "*") {
+		matches, err := filepath.Glob(filename)
+		if err != nil {
+			return []string{}, fmt.Errorf("failed parsing glob filepath\nInputFile=%s\nError=%s", color.Cyan(filename), color.RedBg(err))
+		}
+		return matches, nil
+	} else {
+		return []string{filename}, nil
+	}
+}


Test patch:
diff --git a/example/user_template_test.go b/example/user_template_test.go
new file mode 100644
index 0000000..c206c29
--- /dev/null
+++ b/example/user_template_test.go
@@ -0,0 +1,14 @@
+package example
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestUserTemplateColor(t *testing.T) {
+	assert.Equal(t, OceanColor(0), OceanColorCerulean)
+	assert.Equal(t, true, ParseOceanColorExample())
+	assert.Equal(t, true, ParseOceanColorGlobbedExample())
+	assert.Equal(t, true, ParseOceanColorGlobbedExample2())
+}