/
generator.go
78 lines (65 loc) · 1.69 KB
/
generator.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
package internal
import (
"bytes"
"fmt"
"go/format"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/descriptorpb"
plugin "google.golang.org/protobuf/types/pluginpb"
"path/filepath"
"strings"
)
type Generator interface {
Generate(request *plugin.CodeGeneratorRequest) ([]*plugin.CodeGeneratorResponse_File, error)
}
type generator struct{}
func NewGenerator() Generator {
return new(generator)
}
func (g *generator) Generate(request *plugin.CodeGeneratorRequest) ([]*plugin.CodeGeneratorResponse_File, error) {
var files []*plugin.CodeGeneratorResponse_File
for _, fileToGenerate := range request.FileToGenerate {
var file *descriptorpb.FileDescriptorProto
for _, f := range request.ProtoFile {
if f.GetName() == fileToGenerate {
file = f
}
}
code, err := g.generate(file)
if err != nil {
return nil, err
}
if code == nil {
continue
}
formattedCode, err := format.Source([]byte(*code))
if err != nil {
return nil, err
}
fileName := g.fileName(file)
files = append(files, &plugin.CodeGeneratorResponse_File{
Name: proto.String(fileName),
Content: proto.String(string(formattedCode)),
})
}
return files, nil
}
func (g *generator) generate(file *descriptorpb.FileDescriptorProto) (*string, error) {
if len(file.GetService()) == 0 {
return nil, nil
}
buf := bytes.NewBufferString("")
err := FileTmpl.Execute(buf, file)
if err != nil {
return nil, err
}
out := buf.String()
return &out, nil
}
func (g *generator) fileName(file *descriptorpb.FileDescriptorProto) string {
name := file.GetName()
ext := filepath.Ext(name)
base := strings.TrimSuffix(name, ext)
output := fmt.Sprintf("%s_jgw.pb.go", base)
return output
}