protoc-gen-go-micro/encoder.go

175 lines
4.7 KiB
Go
Raw Normal View History

2016-11-04 23:45:25 +03:00
package main
import (
"bytes"
2016-11-07 10:51:13 +03:00
"log"
2016-11-07 11:23:59 +03:00
"os"
2016-11-06 19:20:57 +03:00
"path/filepath"
2016-11-07 12:27:57 +03:00
"strings"
"text/template"
2016-11-07 11:23:59 +03:00
"time"
2016-11-06 19:20:57 +03:00
2016-11-04 23:45:25 +03:00
"github.com/golang/protobuf/protoc-gen-go/descriptor"
"github.com/golang/protobuf/protoc-gen-go/plugin"
)
type GenericTemplateBasedEncoder struct {
2016-11-06 19:20:57 +03:00
templateDir string
service *descriptor.ServiceDescriptorProto
file *descriptor.FileDescriptorProto
2016-11-07 10:45:02 +03:00
debug bool
2016-11-04 23:45:25 +03:00
}
type Ast struct {
2016-11-07 11:23:59 +03:00
BuildDate time.Time `json:"build-date"`
BuildHostname string `json:"build-hostname"`
BuildUser string `json:"build-user"`
2016-12-14 13:04:34 +03:00
GoPWD string `json:"go-pwd,omitempty"`
2016-11-07 12:27:57 +03:00
PWD string `json:"pwd"`
2016-11-07 11:23:59 +03:00
Debug bool `json:"debug"`
File *descriptor.FileDescriptorProto `json:"file"`
RawFilename string `json:"raw-filename"`
2016-11-07 11:23:59 +03:00
Filename string `json:"filename"`
2016-12-15 12:59:44 +03:00
TemplateDir string `json:"template-dir"`
2016-11-07 11:23:59 +03:00
Service *descriptor.ServiceDescriptorProto `json:"service"`
2016-12-15 13:01:29 +03:00
Environment []string `json:"environment"`
}
2016-11-07 10:45:02 +03:00
func NewGenericTemplateBasedEncoder(templateDir string, service *descriptor.ServiceDescriptorProto, file *descriptor.FileDescriptorProto, debug bool) (e *GenericTemplateBasedEncoder) {
2016-11-04 23:45:25 +03:00
e = &GenericTemplateBasedEncoder{
2016-11-06 19:20:57 +03:00
service: service,
file: file,
templateDir: templateDir,
2016-11-07 10:45:02 +03:00
debug: debug,
2016-11-04 23:45:25 +03:00
}
2016-11-07 10:51:13 +03:00
if debug {
2016-11-09 20:12:00 +03:00
log.Printf("new encoder: file=%q service=%q template-dir=%q", file.GetName(), service.GetName(), templateDir)
2016-11-07 10:51:13 +03:00
}
2016-11-04 23:45:25 +03:00
return
}
2016-11-06 19:20:57 +03:00
func (e *GenericTemplateBasedEncoder) templates() ([]string, error) {
filenames := []string{}
err := filepath.Walk(e.templateDir, func(path string, info os.FileInfo, err error) error {
if err != nil {
return err
2016-11-06 19:20:57 +03:00
}
if info.IsDir() {
return nil
2016-11-06 19:20:57 +03:00
}
if filepath.Ext(path) != ".tmpl" {
return nil
2016-11-06 19:20:57 +03:00
}
rel, err := filepath.Rel(e.templateDir, path)
2016-11-06 19:20:57 +03:00
if err != nil {
return err
2016-11-06 19:20:57 +03:00
}
2016-11-07 10:51:13 +03:00
if e.debug {
log.Printf("new template: %q", rel)
}
2016-11-06 19:20:57 +03:00
filenames = append(filenames, rel)
return nil
})
return filenames, err
2016-11-06 19:20:57 +03:00
}
func (e *GenericTemplateBasedEncoder) genAst(templateFilename string) (*Ast, error) {
// prepare the ast passed to the template engine
2016-11-07 11:23:59 +03:00
hostname, _ := os.Hostname()
2016-11-07 12:27:57 +03:00
pwd, _ := os.Getwd()
goPwd := ""
if os.Getenv("GOPATH") != "" {
goPwd, _ = filepath.Rel(os.Getenv("GOPATH")+"/src", pwd)
if strings.Contains(goPwd, "../") {
goPwd = ""
}
}
ast := Ast{
2016-11-07 11:23:59 +03:00
BuildDate: time.Now(),
BuildHostname: hostname,
BuildUser: os.Getenv("USER"),
2016-11-07 12:27:57 +03:00
PWD: pwd,
GoPWD: goPwd,
2016-11-07 11:23:59 +03:00
File: e.file,
2016-12-15 12:59:44 +03:00
TemplateDir: e.templateDir,
RawFilename: templateFilename,
2016-12-14 13:06:24 +03:00
Filename: "",
2016-12-15 13:01:29 +03:00
Environment: os.Environ(),
2016-12-14 13:06:24 +03:00
Service: e.service,
}
2016-12-14 13:06:24 +03:00
buffer := new(bytes.Buffer)
2016-12-15 14:08:00 +03:00
tmpl, err := template.New("").Funcs(ProtoHelpersFuncMap).Parse(templateFilename)
2016-12-14 13:06:24 +03:00
if err != nil {
return nil, err
}
if err := tmpl.Execute(buffer, ast); err != nil {
return nil, err
}
ast.Filename = buffer.String()
return &ast, nil
}
2016-12-14 13:06:24 +03:00
func (e *GenericTemplateBasedEncoder) buildContent(templateFilename string) (string, string, error) {
// initialize template engine
fullPath := filepath.Join(e.templateDir, templateFilename)
templateName := filepath.Base(fullPath)
2016-12-15 14:08:00 +03:00
tmpl, err := template.New(templateName).Funcs(ProtoHelpersFuncMap).ParseFiles(fullPath)
if err != nil {
2016-12-14 13:06:24 +03:00
return "", "", err
}
ast, err := e.genAst(templateFilename)
if err != nil {
2016-12-14 13:06:24 +03:00
return "", "", err
}
// generate the content
buffer := new(bytes.Buffer)
if err := tmpl.Execute(buffer, ast); err != nil {
2016-12-14 13:06:24 +03:00
return "", "", err
}
2016-12-14 13:06:24 +03:00
return buffer.String(), ast.Filename, nil
}
2016-11-04 23:45:25 +03:00
func (e *GenericTemplateBasedEncoder) Files() []*plugin_go.CodeGeneratorResponse_File {
2016-11-06 19:20:57 +03:00
templates, err := e.templates()
if err != nil {
2016-11-07 10:51:13 +03:00
log.Fatalf("cannot get templates from %q: %v", e.templateDir, err)
2016-11-06 19:20:57 +03:00
}
2016-12-14 13:07:05 +03:00
length := len(templates)
files := make([]*plugin_go.CodeGeneratorResponse_File, 0, length)
errChan := make(chan error, length)
resultChan := make(chan *plugin_go.CodeGeneratorResponse_File, length)
2016-11-06 19:20:57 +03:00
for _, templateFilename := range templates {
2016-12-14 13:07:05 +03:00
go func(tmpl string) {
content, translatedFilename, err := e.buildContent(tmpl)
if err != nil {
errChan <- err
return
}
filename := translatedFilename[:len(translatedFilename)-len(".tmpl")]
resultChan <- &plugin_go.CodeGeneratorResponse_File{
Content: &content,
Name: &filename,
}
}(templateFilename)
}
for i := 0; i < length; i++ {
select {
case f := <-resultChan:
files = append(files, f)
case err = <-errChan:
}
2016-11-04 23:45:25 +03:00
}
2016-12-14 13:07:05 +03:00
if err != nil {
panic(err)
}
2016-11-06 19:20:57 +03:00
return files
2016-11-04 23:45:25 +03:00
}