/
genbm.go
131 lines (116 loc) · 3.16 KB
/
genbm.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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package genbm
import (
"bytes"
"go/format"
"os"
"path/filepath"
"strings"
"github.com/golang/glog"
"github.com/golang/protobuf/protoc-gen-go/descriptor"
plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
"go-common/app/tool/protoc-gen-bm/generator"
)
// New blademaster server code generator
func New(jsonpb bool) generator.Generator {
return &genbm{jsonpb: jsonpb}
}
type genbm struct {
jsonpb bool
}
func (g *genbm) Generate(req *plugin.CodeGeneratorRequest) ([]*plugin.CodeGeneratorResponse_File, error) {
var resp []*plugin.CodeGeneratorResponse_File
files := req.GetProtoFile()
for _, file := range files {
respFile, ok, err := g.generateFile(file)
if err != nil {
return resp, err
}
if ok {
resp = append(resp, respFile)
}
}
return resp, nil
}
func (g *genbm) generateFile(file *descriptor.FileDescriptorProto) (*plugin.CodeGeneratorResponse_File, bool, error) {
glog.V(1).Infof("process proto file %s", file.GetName())
services := file.GetService()
if len(services) == 0 {
glog.V(5).Infof("proto file %s not included service descriptor", file.GetName())
return nil, false, nil
}
var descs []*BMServerDescriptor
for _, service := range services {
server, err := ParseBMServer(service)
if err != nil {
return nil, false, err
}
descs = append(descs, server)
}
buf := new(bytes.Buffer)
goPackageName := GetGoPackageName(file)
gen := NewBMGenerate(goPackageName, descs, g.jsonpb)
if err := gen.Generate(buf); err != nil {
return nil, false, err
}
// format code
data, err := format.Source(buf.Bytes())
if err != nil {
return nil, false, err
}
content := string(data)
// no content
if len(content) == 0 {
return nil, false, nil
}
target := TargetFilePath(file)
glog.V(1).Infof("generate code to %s", target)
return &plugin.CodeGeneratorResponse_File{
Content: &content,
Name: &target,
}, true, nil
}
// TargetFilePath find target file path
func TargetFilePath(file *descriptor.FileDescriptorProto) string {
fpath := file.GetName()
protoDir := filepath.Dir(fpath)
noExt := filepath.Base(fpath)
for i := len(noExt) - 1; i >= 0 && !os.IsPathSeparator(noExt[i]); i-- {
if noExt[i] == '.' {
noExt = noExt[:i]
}
}
target := noExt + ".pb.bm.go"
options := file.GetOptions()
if options != nil {
goPackage := options.GetGoPackage()
if goPackage != "" {
goPackage = strings.Split(goPackage, ";")[0]
if strings.Contains(goPackage, "/") {
return filepath.Join(goPackage, target)
}
}
}
return filepath.Join(protoDir, target)
}
// GetGoPackageName last element from proto package name or go_package option
func GetGoPackageName(file *descriptor.FileDescriptorProto) string {
var goPackageName string
protoPackage := file.GetPackage()
goPackageName = splitLastElem(protoPackage, ".")
options := file.GetOptions()
if options == nil {
return goPackageName
}
if goPackage := options.GetGoPackage(); goPackage != "" {
if strings.Contains(goPackage, ";") {
goPackageName = splitLastElem(goPackage, ";")
} else {
goPackageName = splitLastElem(goPackage, "/")
}
}
return goPackageName
}
func splitLastElem(s string, seq string) string {
seqs := strings.Split(s, seq)
return seqs[len(seqs)-1]
}