-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
221 lines (190 loc) · 6.23 KB
/
main.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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
package main
import (
"flag"
"fmt"
"github.com/firegoblin/gographviz"
"go/ast"
"go/parser"
"go/token"
"os"
"strings"
)
var filename = flag.String("file", "github.com/firegoblin/golangTypeGraph", "file to parse on, relative to $GOPATH/src")
var includeTestFiles = flag.Bool("test", false, "whether or not to include test files in the graph")
var defaultPackageName = flag.String("pkg", "main", "the package that will not have its types prefiexed with package name")
var onlyExports = flag.Bool("exports", false, "marks whether only exported nodes are shown")
var implementMax = flag.Int("imax", 9, "the maximum number of structs implementing an interface before edges are not drawn")
var envVar = flag.String("env", "GOPATH", "environment variable to use instead of GOPATH")
var maxDepth = flag.Int("depth", 1, "maximum depth of recursively searching imports")
var edgelessNodes = flag.Bool("edgeless", true, "include nodes that have no edges connected to them")
var json = flag.Bool("json", false, "include tags for struct fields in print out")
//var operatingSystem = flag.String("os", "linux", "define the os to use when choosing os specific files")
//var operatingArchitecture = flag.String("arch", "amd64", "define the architecture to use when choosing os specific files")
func processTypeDecl(obj *ast.Object, typ *Type, structList *[]*structNode, interfaceList *[]*interfaceNode) {
decl, ok := obj.Decl.(*ast.TypeSpec)
if !ok {
panic("unexpected type in processTypeDecl")
}
node := typ.base.node
switch decl.Type.(type) {
case *ast.InterfaceType:
if node == nil {
*interfaceList = append(*interfaceList, newInterface(decl, typ.base))
} else {
switch n := node.(type) {
case *interfaceNode:
*interfaceList = append(*interfaceList, n.remakeInterface(decl))
case *unknownNode:
*interfaceList = append(*interfaceList, n.remakeInterface(decl))
}
}
//case StructType or redefinied type
default:
if node == nil {
*structList = append(*structList, newStruct(decl, typ.base))
} else {
//fmt.Println(node)
switch n := node.(type) {
case *unknownNode:
*structList = append(*structList, n.remakeStruct(decl))
default:
fmt.Fprintln(os.Stderr, "attempt to recreate struct:", n)
}
}
}
}
func processFuncDecl(decl ast.Decl) {
switch d := decl.(type) {
case *ast.FuncDecl:
f := funcMap.lookupOrAddFromExpr(d.Name.Name, d.Type)
//funcList = append(funcList, f)
if d.Recv != nil {
recv := typeMap.lookupOrAddFromExpr(d.Recv.List[0].Type).base.node
if recv != nil {
//fmt.Println(d.Recv.List[0])
switch r := recv.(type) {
case *structNode:
r.addFunction(f, d.Recv.List[0])
case *unknownNode:
r.addFunction(f, d.Recv.List[0])
default:
panic("trying to add receiver to interface")
}
}
}
}
}
//var osVar = [...]string{"freebsd", "windows", "linux", "dragonfly", "openbsd", "netbsd", "darwin", "plan9", "solaris", "nacl"}
//var arch = [...]string{"386", "amd64", "arm"}
func legalFile(f os.FileInfo) bool {
return *includeTestFiles || !strings.Contains(f.Name(), "_test.go")
}
func containsString(list []string, s string) bool {
for _, v := range list {
if v == s {
return true
}
}
return false
}
func replaceSEL(s string) string {
return strings.Replace(s, "_SEL_", "\\.", -1)
}
func createDOT(interfaceList []*interfaceNode, structList []*structNode) string {
g := gographviz.NewGraph()
g.SetName((*filename)[strings.LastIndex(*filename, "/")+1:])
g.SetDir(true)
for _, s := range structList {
g.AddGraphableNode(g.Name, s)
}
for _, i := range interfaceList {
g.AddGraphableNode(g.Name, i)
}
if !*edgelessNodes {
edges := g.RemoveEdgelessNodes(g.Name)
if len(edges) > 0 {
str := "\"Edgeless nodes removed: "
for _, v := range edges {
str += replaceSEL(v) + ", "
}
str = str[:len(str)-2]
str += "\""
g.AddAttr(g.Name, "label", str)
}
}
return g.String()
}
func main() {
flag.Parse()
fset := token.NewFileSet()
var pkgs map[string]*ast.Package
var err error
gopath := os.Getenv(*envVar) + "/src/"
var structList []*structNode
var interfaceList []*interfaceNode
//var funcList []*function
var directories []string
var depth []int
directories = append(directories, *filename)
depth = append(depth, 0)
var searchedDirectories []string
//loop until directories to search is empty
for len(directories) > 0 {
pkgs, err = parser.ParseDir(fset, gopath+directories[len(directories)-1], legalFile, 0)
//dep is used to
currentDepth := depth[len(depth)-1]
searchedDirectories = append(searchedDirectories, directories[len(directories)-1])
depth = depth[:len(depth)-1]
directories = directories[:len(directories)-1]
//skip this folder if there was an error parsing, usually meaning the directory is not found
if err != nil {
continue
}
for _, pkg := range pkgs {
//remove unexported types/functions if onlyExports
if *onlyExports {
hasExports := ast.PackageExports(pkg)
if !hasExports {
continue
}
}
typeMap.currentPkg = pkg.Name
for _, file := range pkg.Files {
//add imports to directories to check if not at maxDepth yet
if currentDepth < *maxDepth {
for _, impor := range file.Imports {
importName := strings.Trim(impor.Path.Value, "\"")
if !containsString(searchedDirectories, importName) && !containsString(directories, importName) {
depth = append(depth, currentDepth+1)
directories = append(directories, importName)
}
}
}
//add all types to master list before processing delcarations
//minimizes creation of unknown types
for key := range file.Scope.Objects {
typeMap.lookupOrAdd(key)
}
//processes all structs, interfaces, and embedded types
for key, scope := range file.Scope.Objects {
//non-receiver functions are found in scope
if scope.Kind == ast.Typ {
typ := typeMap.lookupOrAdd(key)
processTypeDecl(scope, typ, &structList, &interfaceList)
}
}
//processes all the function declarations
for _, decl := range file.Decls {
processFuncDecl(decl)
}
}
}
}
//set all interfaces' implementedByCache
for _, i := range interfaceList {
i.setImplementedBy(structList)
}
//Create the dot graph and print it out
s := createDOT(interfaceList, structList)
fmt.Println(s)
}