-
Notifications
You must be signed in to change notification settings - Fork 248
/
enumcollector.go
107 lines (101 loc) · 2.62 KB
/
enumcollector.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
package astutils
import (
"go/ast"
"go/parser"
"go/token"
"strings"
)
type EnumCollector struct {
Methods map[string][]MethodMeta
Package PackageMeta
exprString func(ast.Expr) string
Consts map[string][]string
}
// Visit traverse each node from source code
func (sc *EnumCollector) Visit(n ast.Node) ast.Visitor {
return sc.Collect(n)
}
// Collect collects all structs from source code
func (sc *EnumCollector) 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.CONST {
var comments []string
if spec.Doc != nil {
for _, comment := range spec.Doc.List {
comments = append(comments, strings.TrimSpace(strings.TrimPrefix(comment.Text, "//")))
}
}
var typeName string
for _, item := range spec.Specs {
valueSpec := item.(*ast.ValueSpec)
if len(valueSpec.Names) == 0 {
continue
}
switch specType := valueSpec.Type.(type) {
case *ast.Ident:
typeName = specType.Name
case nil:
// useless
//if len(valueSpec.Values) > 0 {
// switch valueExpr := valueSpec.Values[0].(type) {
// case *ast.BasicLit:
// switch valueExpr.Kind {
// case token.INT:
// typeName = "int"
// case token.FLOAT:
// typeName = "float64"
// case token.IMAG:
// typeName = "complex128"
// case token.CHAR:
// typeName = "rune"
// case token.STRING:
// typeName = "string"
// default:
// continue
// }
// }
//}
}
sc.Consts[typeName] = append(sc.Consts[typeName], valueSpec.Names[0].Name)
}
}
}
return nil
}
// NewEnumCollector initializes an EnumCollector
func NewEnumCollector(exprString func(ast.Expr) string) *EnumCollector {
return &EnumCollector{
Methods: make(map[string][]MethodMeta),
Package: PackageMeta{},
exprString: exprString,
Consts: make(map[string][]string),
}
}
func EnumsOf(file string, exprString func(ast.Expr) string) *EnumCollector {
fset := token.NewFileSet()
root, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
if err != nil {
panic(err)
}
sc := NewEnumCollector(exprString)
ast.Walk(sc, root)
return sc
}