forked from josharian/impl
/
impl.go
281 lines (246 loc) · 6.6 KB
/
impl.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
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
// impl generates method stubs for implementing an interface.
package main
import (
"bytes"
"fmt"
"go/ast"
"go/build"
"go/format"
"go/parser"
"go/printer"
"go/token"
"os"
"path/filepath"
"strconv"
"strings"
"text/template"
"code.google.com/p/go.tools/imports"
)
const usage = `impl <recv> <iface>
impl generates method stubs for recv to implement iface.
Examples:
impl 'f *File' io.Reader
impl Murmur hash.Hash
Don't forget the single quotes around the receiver type
to prevent shell globbing.
`
// findInterface returns the import path and identifier of an interface.
// For example, given "http.ResponseWriter", findInterface returns
// "net/http", "ResponseWriter".
func findInterface(iface string) (path string, id string, err error) {
if len(strings.Fields(iface)) != 1 {
return "", "", fmt.Errorf("couldn't parse interface: %s", iface)
}
// Let goimports do the heavy lifting.
src := []byte("package hack\n" + "var i " + iface)
imp, err := imports.Process("", src, nil)
if err != nil {
return "", "", fmt.Errorf("couldn't parse interface: %s", iface)
}
// imp should now contain an appropriate import.
// Parse out the import and the identifier.
fset := token.NewFileSet()
f, err := parser.ParseFile(fset, "", imp, 0)
if err != nil {
panic(err)
}
if len(f.Imports) == 0 {
return "", "", fmt.Errorf("unrecognized interface: %s", iface)
}
raw := f.Imports[0].Path.Value // "io"
path, err = strconv.Unquote(raw) // io
if err != nil {
panic(err)
}
decl := f.Decls[1].(*ast.GenDecl) // var i io.Reader
spec := decl.Specs[0].(*ast.ValueSpec) // i io.Reader
sel := spec.Type.(*ast.SelectorExpr) // io.Reader
id = sel.Sel.Name // Reader
return path, id, nil
}
// Pkg is a parsed build.Package.
type Pkg struct {
*build.Package
*token.FileSet
}
// typeSpec locates the *ast.TypeSpec for type id in the import path.
func typeSpec(path string, id string) (Pkg, *ast.TypeSpec, error) {
pkg, err := build.Import(path, "", 0)
if err != nil {
return Pkg{}, nil, fmt.Errorf("couldn't find package %s: %v", path, err)
}
fset := token.NewFileSet() // share one fset across the whole package
for _, file := range pkg.GoFiles {
f, err := parser.ParseFile(fset, filepath.Join(pkg.Dir, file), nil, 0)
if err != nil {
continue
}
for _, decl := range f.Decls {
decl, ok := decl.(*ast.GenDecl)
if !ok || decl.Tok != token.TYPE {
continue
}
for _, spec := range decl.Specs {
spec := spec.(*ast.TypeSpec)
if spec.Name.Name != id {
continue
}
return Pkg{Package: pkg, FileSet: fset}, spec, nil
}
}
}
return Pkg{}, nil, fmt.Errorf("type %s not found in %s", id, path)
}
// gofmt pretty-prints e.
func (p Pkg) gofmt(e ast.Expr) string {
var buf bytes.Buffer
printer.Fprint(&buf, p.FileSet, e)
return buf.String()
}
// fullType returns the fully qualified type of e.
// Examples, assuming package net/http:
// fullType(int) => "int"
// fullType(Handler) => "http.Handler"
// fullType(io.Reader) => "io.Reader"
// fullType(*Request) => "*http.Request"
func (p Pkg) fullType(e ast.Expr) string {
ast.Inspect(e, func(n ast.Node) bool {
switch n := n.(type) {
case *ast.Ident:
// Using typeSpec instead of IsExported here would be
// more accurate, but it'd be crazy expensive, and if
// the type isn't exported, there's no point trying
// to implement it anyway.
if n.IsExported() {
n.Name = p.Package.Name + "." + n.Name
}
case *ast.SelectorExpr:
return false
}
return true
})
return p.gofmt(e)
}
func (p Pkg) param(field *ast.Field) Param {
var param Param
if len(field.Names) > 0 {
param.Name = field.Names[0].Name
}
param.Type = p.fullType(field.Type)
return param
}
// Method represents a method signature.
type Method struct {
Recv string
Func
}
// Func represents a function signature.
type Func struct {
Name string
Params []Param
Res []Param
}
// Param represents a parameter in a function or method signature.
type Param struct {
Name string
Type string
}
func (p Pkg) funcsig(f *ast.Field) Func {
fn := Func{Name: f.Names[0].Name}
typ := f.Type.(*ast.FuncType)
if typ.Params != nil {
for _, field := range typ.Params.List {
fn.Params = append(fn.Params, p.param(field))
}
}
if typ.Results != nil {
for _, field := range typ.Results.List {
fn.Res = append(fn.Res, p.param(field))
}
}
return fn
}
// funcs returns the set of methods required to implement iface.
// It is called funcs rather than methods because the
// function descriptions are functions; there is no receiver.
func funcs(iface string) ([]Func, error) {
// Locate the interface.
path, id, err := findInterface(iface)
if err != nil {
return nil, err
}
// Parse the package and find the interface declaration.
p, spec, err := typeSpec(path, id)
if err != nil {
return nil, fmt.Errorf("interface %s not found: %s", iface, err)
}
idecl, ok := spec.Type.(*ast.InterfaceType)
if !ok {
return nil, fmt.Errorf("not an interface: %s", iface)
}
if idecl.Methods == nil {
return nil, fmt.Errorf("empty interface: %s", iface)
}
var fns []Func
for _, fndecl := range idecl.Methods.List {
if len(fndecl.Names) == 0 {
// Embedded interface: recurse
embedded, err := funcs(p.fullType(fndecl.Type))
if err != nil {
return nil, err
}
fns = append(fns, embedded...)
continue
}
fn := p.funcsig(fndecl)
fns = append(fns, fn)
}
return fns, nil
}
const stub = "func ({{.Recv}}) {{.Name}}" +
"({{range .Params}}{{.Name}} {{.Type}}, {{end}})" +
"({{range .Res}}{{.Name}} {{.Type}}, {{end}})" +
"{\n}\n\n"
var tmpl = template.Must(template.New("test").Parse(stub))
// genStubs prints nicely formatted method stubs
// for fns using receiver expression recv.
// If recv is not a valid receiver expression,
// genStubs will panic.
func genStubs(recv string, fns []Func) []byte {
var buf bytes.Buffer
for _, fn := range fns {
meth := Method{Recv: recv, Func: fn}
tmpl.Execute(&buf, meth)
}
pretty, err := format.Source(buf.Bytes())
if err != nil {
panic(err)
}
return pretty
}
// validReceiver reports whether recv is a valid receiver expression.
func validReceiver(recv string) bool {
fset := token.NewFileSet()
_, err := parser.ParseFile(fset, "", "package hack\nfunc ("+recv+") Foo()", 0)
return err == nil
}
func main() {
if len(os.Args) != 3 {
fmt.Fprint(os.Stderr, usage)
os.Exit(2)
}
recv, iface := os.Args[1], os.Args[2]
if !validReceiver(recv) {
fatal(fmt.Sprintf("invalid receiver: %q", recv))
}
fns, err := funcs(iface)
if err != nil {
fatal(err)
}
src := genStubs(recv, fns)
fmt.Print(string(src))
}
func fatal(msg interface{}) {
fmt.Fprintln(os.Stderr, msg)
os.Exit(1)
}