Add --template flag to provide filenames of additional templates (#53)

* Add --template flag to provide filenames of additional templates

* Update readme

* Add globbing support to templates

* Share globbing code between --file and --template

Author: charlieparkes

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

example/user_globbed.tmpl

@@ -0,0 +1,3 @@
+func Parse{{.enum.Name}}GlobbedExample() bool {
+    return true
+}

example/user_globbed2.tmpl

@@ -0,0 +1,3 @@
+func Parse{{.enum.Name}}GlobbedExample2() bool {
+    return true
+}

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

example/user_template.tmpl

@@ -0,0 +1,4 @@
+// Additional template
+func Parse{{.enum.Name}}Example() bool {
+    return true
+}

example/user_template_enum.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())
+}

example/user_template_test.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)

generator/generator.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
+	}
+}