protoc-gen-gotemplate/encoder.go

package main

import (
	"bytes"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/golang/protobuf/protoc-gen-go/plugin"
	"github.com/kr/fs"
	"github.com/moul/funcmap"
)

type GenericTemplateBasedEncoder struct {
	templateDir string
	service     *descriptor.ServiceDescriptorProto
	file        *descriptor.FileDescriptorProto
	debug       bool
}

type Ast struct {
	BuildDate     time.Time                          `json:"build-date"`
	BuildHostname string                             `json:"build-hostname"`
	BuildUser     string                             `json:"build-user"`
	GoPWD         string                             `json:"go-pwd",omitempty`
	PWD           string                             `json:"pwd"`
	Debug         bool                               `json:"debug"`
	File          *descriptor.FileDescriptorProto    `json:"file"`
	RawFilename   string                             `json:"raw-filename"`
	Filename      string                             `json:"filename"`
	Service       *descriptor.ServiceDescriptorProto `json:"service"`
}

func NewGenericTemplateBasedEncoder(templateDir string, service *descriptor.ServiceDescriptorProto, file *descriptor.FileDescriptorProto, debug bool) (e *GenericTemplateBasedEncoder) {
	e = &GenericTemplateBasedEncoder{
		service:     service,
		file:        file,
		templateDir: templateDir,
		debug:       debug,
	}

	if debug {
		log.Printf("new encoder: file=%q service=%q template-dir=%q", file.GetName(), service.GetName(), templateDir)
	}

	return
}

func (e *GenericTemplateBasedEncoder) templates() ([]string, error) {
	filenames := []string{}

	walker := fs.Walk(e.templateDir)
	for walker.Step() {
		if err := walker.Err(); err != nil {
			return nil, err
		}

		if walker.Stat().IsDir() {
			continue
		}

		if filepath.Ext(walker.Path()) != ".tmpl" {
			continue
		}

		rel, err := filepath.Rel(e.templateDir, walker.Path())
		if err != nil {
			return nil, err
		}
		if e.debug {
			log.Printf("new template: %q", rel)
		}

		filenames = append(filenames, rel)
	}

	return filenames, nil
}

func (e *GenericTemplateBasedEncoder) genAst(templateFilename string) (*Ast, error) {
	// prepare the ast passed to the template engine
	hostname, _ := os.Hostname()
	pwd, _ := os.Getwd()
	goPwd := ""
	if os.Getenv("GOPATH") != "" {
		goPwd, _ = filepath.Rel(os.Getenv("GOPATH")+"/src", pwd)
		if strings.Contains(goPwd, "../") {
			goPwd = ""
		}
	}
	ast := Ast{
		BuildDate:     time.Now(),
		BuildHostname: hostname,
		BuildUser:     os.Getenv("USER"),
		PWD:           pwd,
		GoPWD:         goPwd,
		File:          e.file,
		RawFilename:   templateFilename,
		// Filename:      "",
		Service: e.service,
	}
	return &ast, nil
}

func (e *GenericTemplateBasedEncoder) buildContent(templateFilename string) (string, error) {
	// initialize template engine
	fullPath := filepath.Join(e.templateDir, templateFilename)
	templateName := filepath.Base(fullPath)
	tmpl, err := template.New(templateName).Funcs(funcmap.FuncMap).ParseFiles(fullPath)
	if err != nil {
		return "", err
	}

	ast, err := e.genAst(templateFilename)
	if err != nil {
		return "", err
	}

	// translate the filename
	ast.Filename, err = e.translateString(templateFilename, templateFilename)
	if err != nil {
		return "", err
	}

	// generate the content
	buffer := new(bytes.Buffer)
	if err := tmpl.Execute(buffer, ast); err != nil {
		return "", err
	}

	return buffer.String(), nil
}

func (e *GenericTemplateBasedEncoder) translateString(input string, templateFilename string) (string, error) {
	buffer := new(bytes.Buffer)
	tmpl, err := template.New("").Funcs(funcmap.FuncMap).Parse(input)
	if err != nil {
		return "", err
	}

	ast, err := e.genAst(templateFilename)
	if err != nil {
		return "", err
	}

	if err := tmpl.Execute(buffer, ast); err != nil {
		return "", err
	}
	return buffer.String(), nil
}

func (e *GenericTemplateBasedEncoder) Files() []*plugin_go.CodeGeneratorResponse_File {
	files := []*plugin_go.CodeGeneratorResponse_File{}

	templates, err := e.templates()
	if err != nil {
		log.Fatalf("cannot get templates from %q: %v", e.templateDir, err)
	}

	for _, templateFilename := range templates {
		content, err := e.buildContent(templateFilename)
		if err != nil {
			panic(err)
		}

		translatedFilename, err := e.translateString(templateFilename, templateFilename)
		if err != nil {
			panic(err)
		}
		filename := translatedFilename[0 : len(translatedFilename)-len(".tmpl")]

		files = append(files, &plugin_go.CodeGeneratorResponse_File{
			Content: &content,
			Name:    &filename,
		})
	}

	return files
}