/
generator.go
114 lines (99 loc) · 2.94 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
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
package retag
import (
"bytes"
"fmt"
"go/ast"
"go/format"
"go/parser"
"go/token"
"google.golang.org/protobuf/compiler/protogen"
"google.golang.org/protobuf/types/pluginpb"
)
var (
fset = token.NewFileSet()
)
type Generator struct {
// Ast goFiles to which this package contains.
goFiles []GoFile
// proto's astFile info
protoFiles []FileInfo
protoGenerator *protogen.Plugin
}
func NewGenerator(protoFiles []FileInfo, protoGenerator *protogen.Plugin) *Generator {
return &Generator{
protoFiles: protoFiles,
protoGenerator: protoGenerator,
}
}
// ParseGoContent analyzes the single package constructed from the patterns and tags.
// ParseGoContent exits if there is an error.
func (g *Generator) ParseGoContent(outerFile *pluginpb.CodeGeneratorResponse_File) {
if outerFile == nil || outerFile.GetContent() == "" {
return
}
const mode = parser.AllErrors | parser.ParseComments
f, err := parser.ParseFile(fset, "", outerFile.GetContent(), mode)
if err != nil {
g.protoGenerator.Error(fmt.Errorf("failed to parse struct tag in field extension: %w", err))
return
}
g.addGoFile(f, outerFile)
}
// addGoFile adds a type checked Package and its syntax goFiles to the generator.
func (g *Generator) addGoFile(astFile *ast.File, outerFile *pluginpb.CodeGeneratorResponse_File) {
g.goFiles = append(g.goFiles, GoFile{
goGenerator: g,
astFile: astFile,
fset: fset,
protoFiles: g.protoFiles,
outerFile: outerFile,
})
}
// Generate produces the rewrite content to proto's Generator.
func (g *Generator) Generate() {
for _, file := range g.goFiles {
// Set the state for this run of the walker.
if file.astFile != nil {
ast.Inspect(file.astFile, file.genDecl)
}
//if file.fileChanged {
// FIXME: always generate *.pb.go, to replace protoc-go, avoid "Tried to write the same file twice"
{
// PrintComment when file is changed by protoc-gen-go-tag.
if file.fileChanged {
PrintComment(file.astFile)
}
var buf bytes.Buffer
err := format.Node(&buf, file.fset, file.astFile)
if err != nil {
g.protoGenerator.Error(fmt.Errorf("failed to format go content: %w", err))
continue
}
// fix Response will always be generated, so add a new generated file directly.
//content := buf.String()
//file.outerFile.Content = &content
_, err = g.protoGenerator.NewGeneratedFile(file.outerFile.GetName(), "").Write(buf.Bytes())
if err != nil {
g.protoGenerator.Error(fmt.Errorf("failed to new generated file to rewrite: %w", err))
continue
}
}
}
}
func PrintComment(file *ast.File) {
if file == nil {
return
}
var list ast.Comment
list.Text = "// Code generated by protoc-gen-go-tag. DO NOT EDIT."
if len(file.Comments) == 0 {
file.Comments = []*ast.CommentGroup{{}}
}
comments := file.Comments[0]
if len(comments.List) == 0 {
comments.List = []*ast.Comment{&list}
}
var lists []*ast.Comment
lists = append(lists, &list)
comments.List = append(lists, comments.List...)
}