-
Notifications
You must be signed in to change notification settings - Fork 400
/
walker.go
136 lines (114 loc) · 2.46 KB
/
walker.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
package mockery
import (
"fmt"
"io"
"io/ioutil"
"os"
"path/filepath"
"regexp"
"strings"
)
type Walker struct {
BaseDir string
Recursive bool
Filter *regexp.Regexp
LimitOne bool
BuildTags []string
}
type WalkerVisitor interface {
VisitWalk(*Interface) error
}
func (this *Walker) Walk(visitor WalkerVisitor) (generated bool) {
parser := NewParser(this.BuildTags)
this.doWalk(parser, this.BaseDir, visitor)
err := parser.Load()
if err != nil {
fmt.Fprintf(os.Stderr, "Error walking: %v\n", err)
os.Exit(1)
}
for _, iface := range parser.Interfaces() {
if !this.Filter.MatchString(iface.Name) {
continue
}
err := visitor.VisitWalk(iface)
if err != nil {
fmt.Fprintf(os.Stderr, "Error walking %s: %s\n", iface.Name, err)
os.Exit(1)
}
generated = true
if this.LimitOne {
return
}
}
return
}
func (this *Walker) doWalk(p *Parser, dir string, visitor WalkerVisitor) (generated bool) {
files, err := ioutil.ReadDir(dir)
if err != nil {
return
}
for _, file := range files {
if strings.HasPrefix(file.Name(), ".") || strings.HasPrefix(file.Name(), "_") {
continue
}
path := filepath.Join(dir, file.Name())
if file.IsDir() {
if this.Recursive {
generated = this.doWalk(p, path, visitor) || generated
if generated && this.LimitOne {
return
}
}
continue
}
if !strings.HasSuffix(path, ".go") || strings.HasSuffix(path, "_test.go") {
continue
}
err = p.Parse(path)
if err != nil {
fmt.Fprintln(os.Stderr, "Error parsing file: ", err)
continue
}
}
return
}
type GeneratorVisitor struct {
InPackage bool
Note string
Osp OutputStreamProvider
// The name of the output package, if InPackage is false (defaults to "mocks")
PackageName string
}
func (this *GeneratorVisitor) VisitWalk(iface *Interface) error {
defer func() {
if r := recover(); r != nil {
fmt.Printf("Unable to generated mock for '%s': %s\n", iface.Name, r)
return
}
}()
var out io.Writer
var pkg string
if this.InPackage {
pkg = filepath.Dir(iface.FileName)
} else {
pkg = this.PackageName
}
out, err, closer := this.Osp.GetWriter(iface)
if err != nil {
fmt.Printf("Unable to get writer for %s: %s", iface.Name, err)
os.Exit(1)
}
defer closer()
gen := NewGenerator(iface, pkg, this.InPackage)
gen.GeneratePrologueNote(this.Note)
gen.GeneratePrologue(pkg)
err = gen.Generate()
if err != nil {
return err
}
err = gen.Write(out)
if err != nil {
return err
}
return nil
}