Skip to content

Commit 6f92a56

Browse files
committed
tweak code
1 parent fb93b42 commit 6f92a56

File tree

3 files changed

+151
-104
lines changed

3 files changed

+151
-104
lines changed

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

+15-104
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
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
// Copyright 2024 CloudWeGo Authors
2+
//
3+
// Licensed to the Apache Software Foundation (ASF) under one or more
4+
// contributor license agreements. See the NOTICE file distributed with
5+
// this work for additional information regarding copyright ownership.
6+
// The ASF licenses this file to You under the Apache License, Version 2.0
7+
// (the "License"); you may not use this file except in compliance with
8+
// the License. You may obtain a copy of the License at
9+
//
10+
// http://www.apache.org/licenses/LICENSE-2.0
11+
//
12+
// Unless required by applicable law or agreed to in writing, software
13+
// distributed under the License is distributed on an "AS IS" BASIS,
14+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
// See the License for the specific language governing permissions and
16+
// limitations under the License.
17+
18+
package generator
19+
20+
import (
21+
"strings"
22+
23+
"google.golang.org/protobuf/compiler/protogen"
24+
)
25+
26+
func generateHessian2Ext(g *generator, gen *protogen.Plugin) error {
27+
for _, file := range gen.Files {
28+
err := g.generateAPIFile(gen, file)
29+
if err != nil {
30+
return err
31+
}
32+
33+
for _, svr := range file.Services {
34+
err := g.generateServiceExt(gen, svr, file)
35+
if err != nil {
36+
return err
37+
}
38+
}
39+
}
40+
return nil
41+
}
42+
43+
// generateAPIFile implements dubbo extension interface every message.
44+
// ref: https://github.com/kitex-contrib/codec-dubbo/blob/dcc4d83669b139d5daa75823d2eda5584d92ebc7/pkg/iface/protocol.go#L22
45+
func (g *generator) generateAPIFile(gen *protogen.Plugin, file *protogen.File) error {
46+
if len(file.Messages) == 0 {
47+
return nil
48+
}
49+
fileName := filePrefix + file.GeneratedFilenamePrefix + "_" + fileExt
50+
genFile := gen.NewGeneratedFile(fileName, "")
51+
request := &ExtFile{
52+
Version: version,
53+
PkgName: file.GoPackageName,
54+
}
55+
if file.Proto != nil && file.Proto.Options != nil {
56+
if file.Proto.Options.GoPackage != nil {
57+
request.GoPackage = *file.Proto.Options.GoPackage
58+
}
59+
if file.Proto.Options.JavaPackage != nil {
60+
request.JavaPackage = *file.Proto.Options.JavaPackage
61+
}
62+
}
63+
messages := file.Messages
64+
for len(messages) != 0 {
65+
var nestedMessages []*protogen.Message
66+
for _, message := range messages {
67+
s := &StructLike{
68+
Name: message.GoIdent.GoName,
69+
GoPackage: request.GoPackage,
70+
JavaPackage: request.JavaPackage,
71+
}
72+
for _, field := range message.Fields {
73+
filed := &Field{
74+
Name: field.GoName,
75+
}
76+
s.Fields = append(s.Fields, filed)
77+
}
78+
request.StructLikes = append(request.StructLikes, s)
79+
nestedMessages = append(nestedMessages, message.Messages...)
80+
}
81+
messages = nestedMessages
82+
}
83+
84+
for _, svr := range file.Services {
85+
s := &ExtService{
86+
Name: svr.GoName,
87+
}
88+
request.Services = append(request.Services, s)
89+
}
90+
91+
var buf strings.Builder
92+
err := g.tpl.ExecuteTemplate(&buf, "structlikes", request)
93+
if err != nil {
94+
return err
95+
}
96+
genFile.P(buf.String())
97+
return nil
98+
}
99+
100+
// generateServiceExt: implements dubbo extension interface for service method args and response
101+
// ref: https://github.com/cloudwego/kitex/blob/8526b3af30fcd321db268cae59a3545a9c6f237f/tool/internal_pkg/generator/generator.go#L392
102+
func (g *generator) generateServiceExt(gen *protogen.Plugin, svr *protogen.Service, file *protogen.File) error {
103+
fileName := filePrefix + string(file.GoImportPath) + "/" + strings.ToLower(svr.GoName) + "/" + fileExt
104+
genFile := gen.NewGeneratedFile(fileName, "")
105+
rSvr := &ExtService{
106+
Version: version,
107+
PkgName: strings.ToLower(svr.GoName),
108+
Name: svr.GoName,
109+
}
110+
for _, method := range svr.Methods {
111+
s := &Function{
112+
Method: method.GoName,
113+
InputType: method.Input.GoIdent.GoName,
114+
OutputType: method.Output.GoIdent.GoName,
115+
}
116+
rSvr.Functions = append(rSvr.Functions, s)
117+
}
118+
var buf strings.Builder
119+
err := g.tpl.ExecuteTemplate(&buf, "service", rSvr)
120+
if err != nil {
121+
return err
122+
}
123+
genFile.P(buf.String())
124+
return nil
125+
}

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

+11
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)