main.go

package main

import (
	"google.golang.org/protobuf/compiler/protogen"
	"strings"
)

func main() {
	protogen.Options{}.Run(func(gen *protogen.Plugin) error {
		for _, f := range gen.Files {
			if !f.Generate {
				continue
			}
			generateFile(gen, f)
		}
		return nil
	})
}

// generateFile generates a _enum.pb.go file containing enum.
func generateFile(gen *protogen.Plugin, file *protogen.File) {

	if len(file.Enums) <= 0 {
		return
	}

	filename := file.GeneratedFilenamePrefix + "_enum.pb.go"
	g := gen.NewGeneratedFile(filename, file.GoImportPath)
	g.P("// Code generated by protoc-gen-go-enum. DO NOT EDIT...")
	g.P()
	g.P("package ", file.GoPackageName)
	g.P()

	for _, enum := range file.Enums {
		// const
		comment := enum.Comments.Leading.String()
		comment = strings.Trim(comment, " ")
		comment = strings.Trim(comment, "\r\n")
		g.P(comment)
		g.P("const (")
		for i, e := range enum.Values {
			comment := e.Comments.Leading.String()
			comment = strings.Trim(comment, " ")
			comment = strings.Trim(comment, "\r\n")
			g.P(e.GoIdent, " int64 ", "=", i, comment)
		}
		g.P(")")

		// map
		g.P("var (")
		g.P(enum.GoIdent, "_TXT=map[int64]string {")
		for _, e := range enum.Values {
			comment := e.Comments.Leading.String()
			comment = strings.TrimPrefix(comment, "//")
			comment = strings.Trim(comment, " ")
			comment = strings.Trim(comment, "\r\n")
			g.P(e.GoIdent, ": \"", comment, "\",")
		}
		g.P("}")
		g.P(")")
	}

}