@@ -20,7 +20,7 @@ import (
20
20
"bytes"
21
21
"encoding/gob"
22
22
"flag"
23
- "io "
23
+ "go/build "
24
24
"io/ioutil"
25
25
"os"
26
26
"os/exec"
@@ -32,69 +32,27 @@ import (
32
32
)
33
33
34
34
var (
35
- progOnly = flag .Bool ("prog_only" , false , "(reflect mode) Only generate the reflection program; write it to stdout." )
35
+ progOnly = flag .Bool ("prog_only" , false , "(reflect mode) Only generate the reflection program; write it to stdout and exit ." )
36
36
execOnly = flag .String ("exec_only" , "" , "(reflect mode) If set, execute this reflection program." )
37
37
buildFlags = flag .String ("build_flags" , "" , "(reflect mode) Additional flags for go build." )
38
38
)
39
39
40
- func Reflect (importPath string , symbols []string ) (* model.Package , error ) {
41
- // TODO: sanity check arguments
42
-
43
- progPath := * execOnly
44
- if * execOnly == "" {
45
- pwd , _ := os .Getwd ()
46
- // We use TempDir instead of TempFile so we can control the filename.
47
- // Try to place the TempDir under pwd, so that if there is some package in
48
- // vendor directory, 'go build' can also load/mock it.
49
- tmpDir , err := ioutil .TempDir (pwd , "gomock_reflect_" )
50
- if err != nil {
51
- return nil , err
52
- }
53
- defer func () { os .RemoveAll (tmpDir ) }()
54
- const progSource = "prog.go"
55
- var progBinary = "prog.bin"
56
- if runtime .GOOS == "windows" {
57
- // Windows won't execute a program unless it has a ".exe" suffix.
58
- progBinary += ".exe"
59
- }
60
-
61
- // Generate program.
62
- var program bytes.Buffer
63
- data := reflectData {
64
- ImportPath : importPath ,
65
- Symbols : symbols ,
66
- }
67
- if err := reflectProgram .Execute (& program , & data ); err != nil {
68
- return nil , err
69
- }
70
- if * progOnly {
71
- io .Copy (os .Stdout , & program )
72
- os .Exit (0 )
73
- }
74
- if err := ioutil .WriteFile (filepath .Join (tmpDir , progSource ), program .Bytes (), 0600 ); err != nil {
75
- return nil , err
76
- }
77
-
78
- cmdArgs := []string {}
79
- cmdArgs = append (cmdArgs , "build" )
80
- if * buildFlags != "" {
81
- cmdArgs = append (cmdArgs , * buildFlags )
82
- }
83
- cmdArgs = append (cmdArgs , "-o" , progBinary , progSource )
84
-
85
- // Build the program.
86
- cmd := exec .Command ("go" , cmdArgs ... )
87
- cmd .Dir = tmpDir
88
- cmd .Stdout = os .Stdout
89
- cmd .Stderr = os .Stderr
90
- if err := cmd .Run (); err != nil {
91
- return nil , err
92
- }
93
- progPath = filepath .Join (tmpDir , progBinary )
40
+ func writeProgram (importPath string , symbols []string ) ([]byte , error ) {
41
+ var program bytes.Buffer
42
+ data := reflectData {
43
+ ImportPath : importPath ,
44
+ Symbols : symbols ,
45
+ }
46
+ if err := reflectProgram .Execute (& program , & data ); err != nil {
47
+ return nil , err
94
48
}
49
+ return program .Bytes (), nil
50
+ }
95
51
96
- // Run it.
97
- cmd := exec .Command (progPath )
52
+ // run the given command and parse the output as a model.Package.
53
+ func run (command string ) (* model.Package , error ) {
54
+ // Run the program.
55
+ cmd := exec .Command (command )
98
56
var stdout bytes.Buffer
99
57
cmd .Stdout = & stdout
100
58
cmd .Stderr = os .Stderr
@@ -110,6 +68,79 @@ func Reflect(importPath string, symbols []string) (*model.Package, error) {
110
68
return & pkg , nil
111
69
}
112
70
71
+ // runInDir writes the given program into the given dir, runs it there, and
72
+ // parses the output as a model.Package.
73
+ func runInDir (program []byte , dir string ) (* model.Package , error ) {
74
+ // We use TempDir instead of TempFile so we can control the filename.
75
+ tmpDir , err := ioutil .TempDir (dir , "gomock_reflect_" )
76
+ if err != nil {
77
+ return nil , err
78
+ }
79
+ defer func () { os .RemoveAll (tmpDir ) }()
80
+ const progSource = "prog.go"
81
+ var progBinary = "prog.bin"
82
+ if runtime .GOOS == "windows" {
83
+ // Windows won't execute a program unless it has a ".exe" suffix.
84
+ progBinary += ".exe"
85
+ }
86
+
87
+ if err := ioutil .WriteFile (filepath .Join (tmpDir , progSource ), program , 0600 ); err != nil {
88
+ return nil , err
89
+ }
90
+
91
+ cmdArgs := []string {}
92
+ cmdArgs = append (cmdArgs , "build" )
93
+ if * buildFlags != "" {
94
+ cmdArgs = append (cmdArgs , * buildFlags )
95
+ }
96
+ cmdArgs = append (cmdArgs , "-o" , progBinary , progSource )
97
+
98
+ // Build the program.
99
+ cmd := exec .Command ("go" , cmdArgs ... )
100
+ cmd .Dir = tmpDir
101
+ cmd .Stdout = os .Stdout
102
+ cmd .Stderr = os .Stderr
103
+ if err := cmd .Run (); err != nil {
104
+ return nil , err
105
+ }
106
+ return run (filepath .Join (tmpDir , progBinary ))
107
+ }
108
+
109
+ func Reflect (importPath string , symbols []string ) (* model.Package , error ) {
110
+ // TODO: sanity check arguments
111
+
112
+ if * execOnly != "" {
113
+ return run (* execOnly )
114
+ }
115
+
116
+ program , err := writeProgram (importPath , symbols )
117
+ if err != nil {
118
+ return nil , err
119
+ }
120
+
121
+ if * progOnly {
122
+ os .Stdout .Write (program )
123
+ os .Exit (0 )
124
+ }
125
+
126
+ wd , _ := os .Getwd ()
127
+
128
+ // Try to run the program in the same directory as the input package.
129
+ if p , err := build .Import (importPath , wd , build .FindOnly ); err == nil {
130
+ dir := p .Dir
131
+ if p , err := runInDir (program , dir ); err == nil {
132
+ return p , nil
133
+ }
134
+ }
135
+
136
+ // Since that didn't work, try to run it in the current working directory.
137
+ if p , err := runInDir (program , wd ); err == nil {
138
+ return p , nil
139
+ }
140
+ // Since that didn't work, try to run it in a standard temp directory.
141
+ return runInDir (program , "" )
142
+ }
143
+
113
144
type reflectData struct {
114
145
ImportPath string
115
146
Symbols []string
0 commit comments