From e1486970f7970b63ae44f04782b68e55ed3d3f24 Mon Sep 17 00:00:00 2001
From: jhayotte <julien.hayotte@gmail.com>
Date: Fri, 12 May 2017 20:46:18 +0200
Subject: [PATCH] add helper getEnumValue

---
 Makefile                              |  1 +
 encoder.go                            |  7 +++++--
 examples/enum/Makefile                | 13 +++++++++++++
 examples/enum/output/enum.txt         |  9 +++++++++
 examples/enum/proto/sample.proto      | 14 ++++++++++++++
 examples/enum/templates/enum.txt.tmpl |  2 ++
 helpers.go                            | 11 +++++++++++
 7 files changed, 55 insertions(+), 2 deletions(-)
 create mode 100644 examples/enum/Makefile
 create mode 100644 examples/enum/output/enum.txt
 create mode 100644 examples/enum/proto/sample.proto
 create mode 100644 examples/enum/templates/enum.txt.tmpl

diff --git a/Makefile b/Makefile
index 7d3f9f6..c9ee6cc 100644
--- a/Makefile
+++ b/Makefile
@@ -8,6 +8,7 @@ install:
 
 .PHONY: test
 test:	install
+	cd examples/enum && make
 	cd examples/import && make
 	cd examples/dummy && make
 	cd examples/flow && make
diff --git a/encoder.go b/encoder.go
index 9de53e9..2444d34 100644
--- a/encoder.go
+++ b/encoder.go
@@ -17,6 +17,7 @@ type GenericTemplateBasedEncoder struct {
 	templateDir    string
 	service        *descriptor.ServiceDescriptorProto
 	file           *descriptor.FileDescriptorProto
+	enum           []*descriptor.EnumDescriptorProto
 	debug          bool
 	destinationDir string
 }
@@ -34,6 +35,7 @@ type Ast struct {
 	Filename       string                             `json:"filename"`
 	TemplateDir    string                             `json:"template-dir"`
 	Service        *descriptor.ServiceDescriptorProto `json:"service"`
+	Enum           []*descriptor.EnumDescriptorProto  `json:"enum"`
 }
 
 func NewGenericServiceTemplateBasedEncoder(templateDir string, service *descriptor.ServiceDescriptorProto, file *descriptor.FileDescriptorProto, debug bool, destinationDir string) (e *GenericTemplateBasedEncoder) {
@@ -43,8 +45,8 @@ func NewGenericServiceTemplateBasedEncoder(templateDir string, service *descript
 		templateDir:    templateDir,
 		debug:          debug,
 		destinationDir: destinationDir,
+		enum:           file.GetEnumType(),
 	}
-
 	if debug {
 		log.Printf("new encoder: file=%q service=%q template-dir=%q", file.GetName(), service.GetName(), templateDir)
 	}
@@ -57,10 +59,10 @@ func NewGenericTemplateBasedEncoder(templateDir string, file *descriptor.FileDes
 		service:        nil,
 		file:           file,
 		templateDir:    templateDir,
+		enum:           file.GetEnumType(),
 		debug:          debug,
 		destinationDir: destinationDir,
 	}
-
 	if debug {
 		log.Printf("new encoder: file=%q template-dir=%q", file.GetName(), templateDir)
 	}
@@ -117,6 +119,7 @@ func (e *GenericTemplateBasedEncoder) genAst(templateFilename string) (*Ast, err
 		RawFilename:    templateFilename,
 		Filename:       "",
 		Service:        e.service,
+		Enum:           e.enum,
 	}
 	buffer := new(bytes.Buffer)
 	tmpl, err := template.New("").Funcs(ProtoHelpersFuncMap).Parse(templateFilename)
diff --git a/examples/enum/Makefile b/examples/enum/Makefile
new file mode 100644
index 0000000..60bc04f
--- /dev/null
+++ b/examples/enum/Makefile
@@ -0,0 +1,13 @@
+.PHONY: build
+build:
+	mkdir -p output
+	protoc -I. --gotemplate_out=template_dir=templates,debug=true,all=true:output proto/*.proto
+
+
+.PHONY: re
+re: clean build
+
+
+.PHONY: clean
+clean:
+	rm -rf output
diff --git a/examples/enum/output/enum.txt b/examples/enum/output/enum.txt
new file mode 100644
index 0000000..4784749
--- /dev/null
+++ b/examples/enum/output/enum.txt
@@ -0,0 +1,9 @@
+-red
+-blue
+-black
+-yellow
+-green
+-dark
+-white
+-gray
+-orange
diff --git a/examples/enum/proto/sample.proto b/examples/enum/proto/sample.proto
new file mode 100644
index 0000000..3d7acd4
--- /dev/null
+++ b/examples/enum/proto/sample.proto
@@ -0,0 +1,14 @@
+syntax = "proto3";
+package Sample;
+
+enum Colors {
+    red = 0;
+    blue = 1;
+    black = 2;
+    yellow = 3;
+    green = 4;
+    dark = 5;
+    white = 6;
+    gray = 7;
+    orange = 8;
+}
\ No newline at end of file
diff --git a/examples/enum/templates/enum.txt.tmpl b/examples/enum/templates/enum.txt.tmpl
new file mode 100644
index 0000000..9bbc9a7
--- /dev/null
+++ b/examples/enum/templates/enum.txt.tmpl
@@ -0,0 +1,2 @@
+{{range $m := "colors" | getEnumValue .Enum }}-{{$m.Name}}
+{{end}}
\ No newline at end of file
diff --git a/helpers.go b/helpers.go
index 4c1a4ab..4d85b1e 100644
--- a/helpers.go
+++ b/helpers.go
@@ -65,6 +65,7 @@ var ProtoHelpersFuncMap = template.FuncMap{
 	},
 	"snakeCase":             xstrings.ToSnakeCase,
 	"getMessageType":        getMessageType,
+	"getEnumValue":          getEnumValue,
 	"isFieldMessage":        isFieldMessage,
 	"isFieldRepeated":       isFieldRepeated,
 	"goType":                goType,
@@ -98,6 +99,16 @@ func getMessageType(f *descriptor.FileDescriptorProto, name string) *descriptor.
 	return nil
 }
 
+func getEnumValue(f []*descriptor.EnumDescriptorProto, name string) []*descriptor.EnumValueDescriptorProto {
+	for _, item := range f {
+		if strings.EqualFold(*item.Name, name) {
+			return item.GetValue()
+		}
+	}
+
+	return nil
+}
+
 func isFieldMessage(f *descriptor.FieldDescriptorProto) bool {
 	if f.Type != nil && *f.Type == descriptor.FieldDescriptorProto_TYPE_MESSAGE {
 		return true