Skip to content

Commit 1551a3d

Browse files
whalecoldwhalecold
committed
tweak code
1 parent fb93b42 commit 1551a3d

File tree

3 files changed

+134
-104
lines changed

3 files changed

+134
-104
lines changed

‎tools/protoc-gen-kitex-dubbo/generator/generator.go‎

Lines changed: 15 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,17 @@ package generator
1919

2020
import (
2121
"html/template"
22-
"strings"
2322

2423
"google.golang.org/protobuf/compiler/protogen"
25-
26-
"github.com/kitex-contrib/codec-dubbo/tools/protoc-gen-kitex-dubbo/templates"
2724
)
2825

26+
var extTemplates []string
27+
28+
// RegisterTemplate add templates to the generator
29+
func RegisterTemplate(t ...string) {
30+
extTemplates = append(extTemplates, t...)
31+
}
32+
2933
const (
3034
filePrefix = "kitex_gen/"
3135
fileExt = "hessian2_ext.go"
@@ -39,128 +43,35 @@ type Generator interface {
3943
type generator struct {
4044
extLang string
4145
tpl *template.Template
46+
funcs []func(g *generator, gen *protogen.Plugin) error
4247
}
4348

4449
func (g *generator) Generate(gen *protogen.Plugin) error {
45-
for _, file := range gen.Files {
46-
err := g.generateAPIFile(gen, file)
50+
for _, f := range g.funcs {
51+
err := f(g, gen)
4752
if err != nil {
4853
return err
4954
}
50-
51-
for _, svr := range file.Services {
52-
err := g.generateServiceExt(gen, svr, file)
53-
if err != nil {
54-
return err
55-
}
56-
}
57-
}
58-
return nil
59-
}
60-
61-
// generateAPIFile implements dubbo extension interface every message.
62-
// ref: https://github.com/kitex-contrib/codec-dubbo/blob/dcc4d83669b139d5daa75823d2eda5584d92ebc7/pkg/iface/protocol.go#L22
63-
func (g *generator) generateAPIFile(gen *protogen.Plugin, file *protogen.File) error {
64-
if len(file.Messages) == 0 {
65-
return nil
66-
}
67-
fileName := filePrefix + file.GeneratedFilenamePrefix + "_" + fileExt
68-
genFile := gen.NewGeneratedFile(fileName, "")
69-
request := &ExtFile{
70-
Version: version,
71-
PkgName: file.GoPackageName,
72-
}
73-
if file.Proto != nil && file.Proto.Options != nil {
74-
if file.Proto.Options.GoPackage != nil {
75-
request.GoPackage = *file.Proto.Options.GoPackage
76-
}
77-
if file.Proto.Options.JavaPackage != nil {
78-
request.JavaPackage = *file.Proto.Options.JavaPackage
79-
}
80-
}
81-
messages := file.Messages
82-
for len(messages) != 0 {
83-
var nestedMessages []*protogen.Message
84-
for _, message := range messages {
85-
s := &StructLike{
86-
Name: message.GoIdent.GoName,
87-
GoPackage: request.GoPackage,
88-
JavaPackage: request.JavaPackage,
89-
}
90-
for _, field := range message.Fields {
91-
filed := &Field{
92-
Name: field.GoName,
93-
}
94-
s.Fields = append(s.Fields, filed)
95-
}
96-
request.StructLikes = append(request.StructLikes, s)
97-
nestedMessages = append(nestedMessages, message.Messages...)
98-
}
99-
messages = nestedMessages
100-
}
101-
102-
for _, svr := range file.Services {
103-
s := &ExtService{
104-
Name: svr.GoName,
105-
}
106-
request.Services = append(request.Services, s)
10755
}
108-
109-
var buf strings.Builder
110-
err := g.tpl.ExecuteTemplate(&buf, "structlikes", request)
111-
if err != nil {
112-
return err
113-
}
114-
genFile.P(buf.String())
115-
return nil
116-
}
117-
118-
// generateServiceExt: implements dubbo extension interface for service method args and response
119-
// ref: https://github.com/cloudwego/kitex/blob/8526b3af30fcd321db268cae59a3545a9c6f237f/tool/internal_pkg/generator/generator.go#L392
120-
func (g *generator) generateServiceExt(gen *protogen.Plugin, svr *protogen.Service, file *protogen.File) error {
121-
fileName := filePrefix + string(file.GoImportPath) + "/" + strings.ToLower(svr.GoName) + "/" + fileExt
122-
genFile := gen.NewGeneratedFile(fileName, "")
123-
rSvr := &ExtService{
124-
Version: version,
125-
PkgName: strings.ToLower(svr.GoName),
126-
Name: svr.GoName,
127-
}
128-
for _, method := range svr.Methods {
129-
s := &Function{
130-
Method: method.GoName,
131-
InputType: method.Input.GoIdent.GoName,
132-
OutputType: method.Output.GoIdent.GoName,
133-
}
134-
rSvr.Functions = append(rSvr.Functions, s)
135-
}
136-
var buf strings.Builder
137-
err := g.tpl.ExecuteTemplate(&buf, "service", rSvr)
138-
if err != nil {
139-
return err
140-
}
141-
genFile.P(buf.String())
14256
return nil
14357
}
14458

14559
// New implements Generator
14660
func New(req *protogen.Plugin, lang string) (Generator, error) {
14761
tpl := template.New("kitex-dubbo")
148-
allTemplates := []string{
149-
templates.StructLikes,
150-
templates.Header,
151-
templates.StructLike,
152-
templates.Service,
153-
templates.JavaClassName,
154-
}
15562
var err error
156-
for _, temp := range allTemplates {
63+
for _, temp := range extTemplates {
15764
tpl, err = tpl.Parse(temp)
15865
if err != nil {
15966
return nil, err
16067
}
16168
}
69+
funcs := []func(g *generator, gen *protogen.Plugin) error{
70+
generateHessian2Ext,
71+
}
16272
return &generator{
16373
tpl: tpl,
16474
extLang: lang,
75+
funcs: funcs,
16576
}, nil
16677
}

‎tools/protoc-gen-kitex-dubbo/generator/hessian2.go‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
package generator
2+
3+
import (
4+
"strings"
5+
6+
"google.golang.org/protobuf/compiler/protogen"
7+
)
8+
9+
func generateHessian2Ext(g *generator, gen *protogen.Plugin) error {
10+
for _, file := range gen.Files {
11+
err := g.generateAPIFile(gen, file)
12+
if err != nil {
13+
return err
14+
}
15+
16+
for _, svr := range file.Services {
17+
err := g.generateServiceExt(gen, svr, file)
18+
if err != nil {
19+
return err
20+
}
21+
}
22+
}
23+
return nil
24+
}
25+
26+
// generateAPIFile implements dubbo extension interface every message.
27+
// ref: https://github.com/kitex-contrib/codec-dubbo/blob/dcc4d83669b139d5daa75823d2eda5584d92ebc7/pkg/iface/protocol.go#L22
28+
func (g *generator) generateAPIFile(gen *protogen.Plugin, file *protogen.File) error {
29+
if len(file.Messages) == 0 {
30+
return nil
31+
}
32+
fileName := filePrefix + file.GeneratedFilenamePrefix + "_" + fileExt
33+
genFile := gen.NewGeneratedFile(fileName, "")
34+
request := &ExtFile{
35+
Version: version,
36+
PkgName: file.GoPackageName,
37+
}
38+
if file.Proto != nil && file.Proto.Options != nil {
39+
if file.Proto.Options.GoPackage != nil {
40+
request.GoPackage = *file.Proto.Options.GoPackage
41+
}
42+
if file.Proto.Options.JavaPackage != nil {
43+
request.JavaPackage = *file.Proto.Options.JavaPackage
44+
}
45+
}
46+
messages := file.Messages
47+
for len(messages) != 0 {
48+
var nestedMessages []*protogen.Message
49+
for _, message := range messages {
50+
s := &StructLike{
51+
Name: message.GoIdent.GoName,
52+
GoPackage: request.GoPackage,
53+
JavaPackage: request.JavaPackage,
54+
}
55+
for _, field := range message.Fields {
56+
filed := &Field{
57+
Name: field.GoName,
58+
}
59+
s.Fields = append(s.Fields, filed)
60+
}
61+
request.StructLikes = append(request.StructLikes, s)
62+
nestedMessages = append(nestedMessages, message.Messages...)
63+
}
64+
messages = nestedMessages
65+
}
66+
67+
for _, svr := range file.Services {
68+
s := &ExtService{
69+
Name: svr.GoName,
70+
}
71+
request.Services = append(request.Services, s)
72+
}
73+
74+
var buf strings.Builder
75+
err := g.tpl.ExecuteTemplate(&buf, "structlikes", request)
76+
if err != nil {
77+
return err
78+
}
79+
genFile.P(buf.String())
80+
return nil
81+
}
82+
83+
// generateServiceExt: implements dubbo extension interface for service method args and response
84+
// ref: https://github.com/cloudwego/kitex/blob/8526b3af30fcd321db268cae59a3545a9c6f237f/tool/internal_pkg/generator/generator.go#L392
85+
func (g *generator) generateServiceExt(gen *protogen.Plugin, svr *protogen.Service, file *protogen.File) error {
86+
fileName := filePrefix + string(file.GoImportPath) + "/" + strings.ToLower(svr.GoName) + "/" + fileExt
87+
genFile := gen.NewGeneratedFile(fileName, "")
88+
rSvr := &ExtService{
89+
Version: version,
90+
PkgName: strings.ToLower(svr.GoName),
91+
Name: svr.GoName,
92+
}
93+
for _, method := range svr.Methods {
94+
s := &Function{
95+
Method: method.GoName,
96+
InputType: method.Input.GoIdent.GoName,
97+
OutputType: method.Output.GoIdent.GoName,
98+
}
99+
rSvr.Functions = append(rSvr.Functions, s)
100+
}
101+
var buf strings.Builder
102+
err := g.tpl.ExecuteTemplate(&buf, "service", rSvr)
103+
if err != nil {
104+
return err
105+
}
106+
genFile.P(buf.String())
107+
return nil
108+
}

‎tools/protoc-gen-kitex-dubbo/templates/hessian_extension.go‎

Lines changed: 11 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)