@@ -19,13 +19,17 @@ package generator
19
19
20
20
import (
21
21
"html/template"
22
- "strings"
23
22
24
23
"google.golang.org/protobuf/compiler/protogen"
25
-
26
- "github.com/kitex-contrib/codec-dubbo/tools/protoc-gen-kitex-dubbo/templates"
27
24
)
28
25
26
+ var extTemplates []string
27
+
28
+ // RegisterTemplate add templates to the generator
29
+ func RegisterTemplate (t ... string ) {
30
+ extTemplates = append (extTemplates , t ... )
31
+ }
32
+
29
33
const (
30
34
filePrefix = "kitex_gen/"
31
35
fileExt = "hessian2_ext.go"
@@ -39,128 +43,35 @@ type Generator interface {
39
43
type generator struct {
40
44
extLang string
41
45
tpl * template.Template
46
+ funcs []func (g * generator , gen * protogen.Plugin ) error
42
47
}
43
48
44
49
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 )
47
52
if err != nil {
48
53
return err
49
54
}
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 )
107
55
}
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 ())
142
56
return nil
143
57
}
144
58
145
59
// New implements Generator
146
60
func New (req * protogen.Plugin , lang string ) (Generator , error ) {
147
61
tpl := template .New ("kitex-dubbo" )
148
- allTemplates := []string {
149
- templates .StructLikes ,
150
- templates .Header ,
151
- templates .StructLike ,
152
- templates .Service ,
153
- templates .JavaClassName ,
154
- }
155
62
var err error
156
- for _ , temp := range allTemplates {
63
+ for _ , temp := range extTemplates {
157
64
tpl , err = tpl .Parse (temp )
158
65
if err != nil {
159
66
return nil , err
160
67
}
161
68
}
69
+ funcs := []func (g * generator , gen * protogen.Plugin ) error {
70
+ generateHessian2Ext ,
71
+ }
162
72
return & generator {
163
73
tpl : tpl ,
164
74
extLang : lang ,
75
+ funcs : funcs ,
165
76
}, nil
166
77
}
0 commit comments