-
Notifications
You must be signed in to change notification settings - Fork 45
/
generater.go
129 lines (107 loc) · 2.27 KB
/
generater.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
package enum
import (
"bytes"
"fmt"
"go/ast"
"go/format"
"go/parser"
"go/token"
"os"
"text/template"
"github.com/pkg/errors"
)
// G Generater
var G = Generater{
t: template.New("enum"),
}
// Generater 用于生成枚举的生成器
type Generater struct {
t *template.Template
file string
Marshal bool
ProtobufExt bool
}
// NewRenderParams todo
func NewRenderParams() *RenderParams {
return &RenderParams{
Enums: NewEnumSet(),
Stringer: true,
ValueMap: true,
Backquote: "`",
}
}
// RenderParams 模板渲染需要的参数
type RenderParams struct {
PKG string
Backquote string
Enums *Set
Marshal bool
Stringer bool
ValueMap bool
}
// Generate 生成文件
func (g *Generater) Generate(file string) ([]byte, error) {
g.file = file
params, err := g.parse()
if err != nil {
return nil, err
}
params.Marshal = g.Marshal
if g.ProtobufExt {
params.Stringer = false
params.ValueMap = false
}
if params.Enums.Length() == 0 {
return []byte{}, nil
}
return g.gen(params)
}
// 兼容命令行测试
func (g *Generater) getFile() string {
if g.file != "" {
return g.file
}
return os.Getenv("GOFILE")
}
// 解析代码源文件,获取常量和类型
func (g *Generater) parse() (*RenderParams, error) {
fset := token.NewFileSet()
f, err := parser.ParseFile(fset, g.getFile(), nil, parser.ParseComments)
if err != nil {
return nil, fmt.Errorf("parse file error, %s", err)
}
params := NewRenderParams()
params.PKG = f.Name.Name
for _, decl := range f.Decls {
switch d := decl.(type) {
case *ast.GenDecl:
switch d.Tok {
case token.CONST:
for _, spec := range d.Specs {
vs, _ := spec.(*ast.ValueSpec)
ident := vs.Names[0]
doc := vs.Doc.Text()
var enum *Enum
vst, _ := vs.Type.(*ast.Ident)
if vst != nil {
enum = params.Enums.Get(vst.Name)
enum.Add(NewItem(ident.Name, doc))
}
}
}
}
}
return params, nil
}
func (g *Generater) gen(params *RenderParams) ([]byte, error) {
buf := bytes.NewBufferString("")
t, err := g.t.Parse(tmp)
if err != nil {
return nil, errors.Wrapf(err, "template init err")
}
err = t.Execute(buf, params)
if err != nil {
return nil, errors.Wrapf(err, "template data err")
}
return format.Source(buf.Bytes())
}