-
Notifications
You must be signed in to change notification settings - Fork 249
/
structcollector.go
170 lines (158 loc) · 4.67 KB
/
structcollector.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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
package astutils
import (
"github.com/sirupsen/logrus"
"go/ast"
"go/parser"
"go/token"
"regexp"
"strings"
"unicode"
)
// StructCollector collect structs by parsing source code
type StructCollector struct {
Structs []StructMeta
Methods map[string][]MethodMeta
Package PackageMeta
NonStructTypeMap map[string]ast.Expr
exprString func(ast.Expr) string
enums map[string]EnumMeta
}
// Visit traverse each node from source code
func (sc *StructCollector) Visit(n ast.Node) ast.Visitor {
return sc.Collect(n)
}
// Collect collects all structs from source code
func (sc *StructCollector) Collect(n ast.Node) ast.Visitor {
switch spec := n.(type) {
case *ast.Package:
return sc
case *ast.File: // actually it is package name
sc.Package = PackageMeta{
Name: spec.Name.Name,
}
return sc
case *ast.FuncDecl:
if spec.Recv != nil {
typeName := strings.TrimPrefix(sc.exprString(spec.Recv.List[0].Type), "*")
methods, _ := sc.Methods[typeName]
methods = append(methods, GetMethodMeta(spec))
if sc.Methods == nil {
sc.Methods = make(map[string][]MethodMeta)
}
sc.Methods[typeName] = methods
}
case *ast.GenDecl:
if spec.Tok == token.TYPE {
var comments []string
if spec.Doc != nil {
for _, comment := range spec.Doc.List {
comments = append(comments, strings.TrimSpace(strings.TrimPrefix(comment.Text, "//")))
}
}
for _, item := range spec.Specs {
typeSpec := item.(*ast.TypeSpec)
typeName := typeSpec.Name.Name
logrus.Printf("Type: name=%s\n", typeName)
switch specType := typeSpec.Type.(type) {
case *ast.StructType:
structmeta := NewStructMeta(specType, sc.exprString)
structmeta.Name = typeName
structmeta.Comments = comments
structmeta.IsExport = unicode.IsUpper(rune(typeName[0]))
sc.Structs = append(sc.Structs, structmeta)
default:
sc.NonStructTypeMap[typeName] = typeSpec.Type
}
}
}
}
return nil
}
// DocFlatEmbed flatten embed struct fields
func (sc *StructCollector) DocFlatEmbed() []StructMeta {
structMap := make(map[string]StructMeta)
for _, structMeta := range sc.Structs {
if _, exists := structMap[structMeta.Name]; !exists {
structMap[structMeta.Name] = structMeta
}
}
var exStructs []StructMeta
for _, structMeta := range sc.Structs {
if !structMeta.IsExport {
continue
}
exStructs = append(exStructs, structMeta)
}
re := regexp.MustCompile(`json:"(.*?)"`)
var result []StructMeta
for _, structMeta := range exStructs {
_structMeta := StructMeta{
Name: structMeta.Name,
Fields: make([]FieldMeta, 0),
Comments: make([]string, len(structMeta.Comments)),
IsExport: true,
}
copy(_structMeta.Comments, structMeta.Comments)
fieldMap := make(map[string]FieldMeta)
embedFieldMap := make(map[string]FieldMeta)
for _, fieldMeta := range structMeta.Fields {
if strings.HasPrefix(fieldMeta.Type, "embed") {
if re.MatchString(fieldMeta.Tag) {
fieldMeta.Type = strings.TrimPrefix(fieldMeta.Type, "embed:")
_structMeta.Fields = append(_structMeta.Fields, fieldMeta)
fieldMap[fieldMeta.Name] = fieldMeta
} else {
if embedded, exists := structMap[fieldMeta.Name]; exists {
for _, field := range embedded.Fields {
if !field.IsExport {
continue
}
embedFieldMap[field.Name] = field
}
}
}
} else if fieldMeta.IsExport {
_structMeta.Fields = append(_structMeta.Fields, fieldMeta)
fieldMap[fieldMeta.Name] = fieldMeta
}
}
for key, field := range embedFieldMap {
if _, exists := fieldMap[key]; !exists {
_structMeta.Fields = append(_structMeta.Fields, field)
}
}
result = append(result, _structMeta)
}
return result
}
type StructCollectorOption func(collector *StructCollector)
func WithEnums(enums map[string]EnumMeta) StructCollectorOption {
return func(collector *StructCollector) {
collector.enums = enums
}
}
// NewStructCollector initializes an StructCollector
func NewStructCollector(exprString func(ast.Expr) string, opts ...StructCollectorOption) *StructCollector {
sc := &StructCollector{
Structs: nil,
Methods: make(map[string][]MethodMeta),
Package: PackageMeta{},
NonStructTypeMap: make(map[string]ast.Expr),
exprString: exprString,
}
for _, opt := range opts {
opt(sc)
}
return sc
}
// BuildStructCollector initializes an StructCollector and collects structs
func BuildStructCollector(file string, exprString func(ast.Expr) string) StructCollector {
sc := NewStructCollector(exprString)
fset := token.NewFileSet()
root, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
if err != nil {
logrus.Panicln(err)
}
ast.Walk(sc, root)
return *sc
}