1
2
3
4
5
6 package astutil
7
8 import (
9 "fmt"
10 "go/ast"
11 "go/token"
12 "strconv"
13 "strings"
14 )
15
16
17 func AddImport(fset *token.FileSet, f *ast.File, path string) (added bool) {
18 return AddNamedImport(fset, f, "", path)
19 }
20
21
22
23
24
25
26
27
28
29
30
31 func AddNamedImport(fset *token.FileSet, f *ast.File, name, path string) (added bool) {
32 if imports(f, name, path) {
33 return false
34 }
35
36 newImport := &ast.ImportSpec{
37 Path: &ast.BasicLit{
38 Kind: token.STRING,
39 Value: strconv.Quote(path),
40 },
41 }
42 if name != "" {
43 newImport.Name = &ast.Ident{Name: name}
44 }
45
46
47
48
49
50 var (
51 bestMatch = -1
52 lastImport = -1
53 impDecl *ast.GenDecl
54 impIndex = -1
55
56 isThirdPartyPath = isThirdParty(path)
57 )
58 for i, decl := range f.Decls {
59 gen, ok := decl.(*ast.GenDecl)
60 if ok && gen.Tok == token.IMPORT {
61 lastImport = i
62
63
64 if declImports(gen, "C") {
65 continue
66 }
67
68
69 if len(gen.Specs) == 0 && bestMatch == -1 {
70 impDecl = gen
71 }
72
73
74
75
76
77
78
79
80
81
82
83 seenAnyThirdParty := false
84 for j, spec := range gen.Specs {
85 impspec := spec.(*ast.ImportSpec)
86 p := importPath(impspec)
87 n := matchLen(p, path)
88 if n > bestMatch || (bestMatch == 0 && !seenAnyThirdParty && isThirdPartyPath) {
89 bestMatch = n
90 impDecl = gen
91 impIndex = j
92 }
93 seenAnyThirdParty = seenAnyThirdParty || isThirdParty(p)
94 }
95 }
96 }
97
98
99 if impDecl == nil {
100 impDecl = &ast.GenDecl{
101 Tok: token.IMPORT,
102 }
103 if lastImport >= 0 {
104 impDecl.TokPos = f.Decls[lastImport].End()
105 } else {
106
107
108
109
110 impDecl.TokPos = f.Package
111
112 file := fset.File(f.Package)
113 pkgLine := file.Line(f.Package)
114 for _, c := range f.Comments {
115 if file.Line(c.Pos()) > pkgLine {
116 break
117 }
118
119 impDecl.TokPos = c.End() + 2
120 }
121 }
122 f.Decls = append(f.Decls, nil)
123 copy(f.Decls[lastImport+2:], f.Decls[lastImport+1:])
124 f.Decls[lastImport+1] = impDecl
125 }
126
127
128 insertAt := 0
129 if impIndex >= 0 {
130
131 insertAt = impIndex + 1
132 }
133 impDecl.Specs = append(impDecl.Specs, nil)
134 copy(impDecl.Specs[insertAt+1:], impDecl.Specs[insertAt:])
135 impDecl.Specs[insertAt] = newImport
136 pos := impDecl.Pos()
137 if insertAt > 0 {
138
139
140 if spec, ok := impDecl.Specs[insertAt-1].(*ast.ImportSpec); ok && spec.Comment != nil {
141 pos = spec.Comment.End()
142 } else {
143
144
145 pos = impDecl.Specs[insertAt-1].Pos()
146 }
147 }
148 if newImport.Name != nil {
149 newImport.Name.NamePos = pos
150 }
151 newImport.Path.ValuePos = pos
152 newImport.EndPos = pos
153
154
155 if len(impDecl.Specs) == 1 {
156
157 impDecl.Lparen = token.NoPos
158 } else if !impDecl.Lparen.IsValid() {
159
160 impDecl.Lparen = impDecl.Specs[0].Pos()
161 }
162
163 f.Imports = append(f.Imports, newImport)
164
165 if len(f.Decls) <= 1 {
166 return true
167 }
168
169
170 var first *ast.GenDecl
171 for i := 0; i < len(f.Decls); i++ {
172 decl := f.Decls[i]
173 gen, ok := decl.(*ast.GenDecl)
174 if !ok || gen.Tok != token.IMPORT || declImports(gen, "C") {
175 continue
176 }
177 if first == nil {
178 first = gen
179 continue
180 }
181
182
183 first.Lparen = first.Pos()
184
185 for _, spec := range gen.Specs {
186 spec.(*ast.ImportSpec).Path.ValuePos = first.Pos()
187 first.Specs = append(first.Specs, spec)
188 }
189 f.Decls = append(f.Decls[:i], f.Decls[i+1:]...)
190 i--
191 }
192
193 return true
194 }
195
196 func isThirdParty(importPath string) bool {
197
198
199 return strings.Contains(importPath, ".")
200 }
201
202
203
204 func DeleteImport(fset *token.FileSet, f *ast.File, path string) (deleted bool) {
205 return DeleteNamedImport(fset, f, "", path)
206 }
207
208
209
210 func DeleteNamedImport(fset *token.FileSet, f *ast.File, name, path string) (deleted bool) {
211 var delspecs []*ast.ImportSpec
212 var delcomments []*ast.CommentGroup
213
214
215 for i := 0; i < len(f.Decls); i++ {
216 decl := f.Decls[i]
217 gen, ok := decl.(*ast.GenDecl)
218 if !ok || gen.Tok != token.IMPORT {
219 continue
220 }
221 for j := 0; j < len(gen.Specs); j++ {
222 spec := gen.Specs[j]
223 impspec := spec.(*ast.ImportSpec)
224 if importName(impspec) != name || importPath(impspec) != path {
225 continue
226 }
227
228
229
230 delspecs = append(delspecs, impspec)
231 deleted = true
232 copy(gen.Specs[j:], gen.Specs[j+1:])
233 gen.Specs = gen.Specs[:len(gen.Specs)-1]
234
235
236
237 if len(gen.Specs) == 0 {
238 copy(f.Decls[i:], f.Decls[i+1:])
239 f.Decls = f.Decls[:len(f.Decls)-1]
240 i--
241 break
242 } else if len(gen.Specs) == 1 {
243 if impspec.Doc != nil {
244 delcomments = append(delcomments, impspec.Doc)
245 }
246 if impspec.Comment != nil {
247 delcomments = append(delcomments, impspec.Comment)
248 }
249 for _, cg := range f.Comments {
250
251 if cg.End() < impspec.Pos() && fset.Position(cg.End()).Line == fset.Position(impspec.Pos()).Line {
252 delcomments = append(delcomments, cg)
253 break
254 }
255 }
256
257 spec := gen.Specs[0].(*ast.ImportSpec)
258
259
260 if spec.Doc != nil {
261 for fset.Position(gen.TokPos).Line+1 < fset.Position(spec.Doc.Pos()).Line {
262 fset.File(gen.TokPos).MergeLine(fset.Position(gen.TokPos).Line)
263 }
264 }
265 for _, cg := range f.Comments {
266 if cg.End() < spec.Pos() && fset.Position(cg.End()).Line == fset.Position(spec.Pos()).Line {
267 for fset.Position(gen.TokPos).Line+1 < fset.Position(spec.Pos()).Line {
268 fset.File(gen.TokPos).MergeLine(fset.Position(gen.TokPos).Line)
269 }
270 break
271 }
272 }
273 }
274 if j > 0 {
275 lastImpspec := gen.Specs[j-1].(*ast.ImportSpec)
276 lastLine := fset.PositionFor(lastImpspec.Path.ValuePos, false).Line
277 line := fset.PositionFor(impspec.Path.ValuePos, false).Line
278
279
280
281 if line-lastLine > 1 || !gen.Rparen.IsValid() {
282
283
284
285
286 } else if line != fset.File(gen.Rparen).LineCount() {
287
288 fset.File(gen.Rparen).MergeLine(line)
289 }
290 }
291 j--
292 }
293 }
294
295
296 for i := 0; i < len(f.Imports); i++ {
297 imp := f.Imports[i]
298 for j, del := range delspecs {
299 if imp == del {
300 copy(f.Imports[i:], f.Imports[i+1:])
301 f.Imports = f.Imports[:len(f.Imports)-1]
302 copy(delspecs[j:], delspecs[j+1:])
303 delspecs = delspecs[:len(delspecs)-1]
304 i--
305 break
306 }
307 }
308 }
309
310
311 for i := 0; i < len(f.Comments); i++ {
312 cg := f.Comments[i]
313 for j, del := range delcomments {
314 if cg == del {
315 copy(f.Comments[i:], f.Comments[i+1:])
316 f.Comments = f.Comments[:len(f.Comments)-1]
317 copy(delcomments[j:], delcomments[j+1:])
318 delcomments = delcomments[:len(delcomments)-1]
319 i--
320 break
321 }
322 }
323 }
324
325 if len(delspecs) > 0 {
326 panic(fmt.Sprintf("deleted specs from Decls but not Imports: %v", delspecs))
327 }
328
329 return
330 }
331
332
333 func RewriteImport(fset *token.FileSet, f *ast.File, oldPath, newPath string) (rewrote bool) {
334 for _, imp := range f.Imports {
335 if importPath(imp) == oldPath {
336 rewrote = true
337
338
339 imp.EndPos = imp.End()
340 imp.Path.Value = strconv.Quote(newPath)
341 }
342 }
343 return
344 }
345
346
347 func UsesImport(f *ast.File, path string) (used bool) {
348 spec := importSpec(f, path)
349 if spec == nil {
350 return
351 }
352
353 name := spec.Name.String()
354 switch name {
355 case "<nil>":
356
357
358 lastSlash := strings.LastIndex(path, "/")
359 if lastSlash == -1 {
360 name = path
361 } else {
362 name = path[lastSlash+1:]
363 }
364 case "_", ".":
365
366 return true
367 }
368
369 ast.Walk(visitFn(func(n ast.Node) {
370 sel, ok := n.(*ast.SelectorExpr)
371 if ok && isTopName(sel.X, name) {
372 used = true
373 }
374 }), f)
375
376 return
377 }
378
379 type visitFn func(node ast.Node)
380
381 func (fn visitFn) Visit(node ast.Node) ast.Visitor {
382 fn(node)
383 return fn
384 }
385
386
387 func imports(f *ast.File, name, path string) bool {
388 for _, s := range f.Imports {
389 if importName(s) == name && importPath(s) == path {
390 return true
391 }
392 }
393 return false
394 }
395
396
397
398 func importSpec(f *ast.File, path string) *ast.ImportSpec {
399 for _, s := range f.Imports {
400 if importPath(s) == path {
401 return s
402 }
403 }
404 return nil
405 }
406
407
408
409 func importName(s *ast.ImportSpec) string {
410 if s.Name == nil {
411 return ""
412 }
413 return s.Name.Name
414 }
415
416
417
418 func importPath(s *ast.ImportSpec) string {
419 t, err := strconv.Unquote(s.Path.Value)
420 if err != nil {
421 return ""
422 }
423 return t
424 }
425
426
427 func declImports(gen *ast.GenDecl, path string) bool {
428 if gen.Tok != token.IMPORT {
429 return false
430 }
431 for _, spec := range gen.Specs {
432 impspec := spec.(*ast.ImportSpec)
433 if importPath(impspec) == path {
434 return true
435 }
436 }
437 return false
438 }
439
440
441 func matchLen(x, y string) int {
442 n := 0
443 for i := 0; i < len(x) && i < len(y) && x[i] == y[i]; i++ {
444 if x[i] == '/' {
445 n++
446 }
447 }
448 return n
449 }
450
451
452 func isTopName(n ast.Expr, name string) bool {
453 id, ok := n.(*ast.Ident)
454 return ok && id.Name == name && id.Obj == nil
455 }
456
457
458 func Imports(fset *token.FileSet, f *ast.File) [][]*ast.ImportSpec {
459 var groups [][]*ast.ImportSpec
460
461 for _, decl := range f.Decls {
462 genDecl, ok := decl.(*ast.GenDecl)
463 if !ok || genDecl.Tok != token.IMPORT {
464 break
465 }
466
467 group := []*ast.ImportSpec{}
468
469 var lastLine int
470 for _, spec := range genDecl.Specs {
471 importSpec := spec.(*ast.ImportSpec)
472 pos := importSpec.Path.ValuePos
473 line := fset.Position(pos).Line
474 if lastLine > 0 && pos > 0 && line-lastLine > 1 {
475 groups = append(groups, group)
476 group = []*ast.ImportSpec{}
477 }
478 group = append(group, importSpec)
479 lastLine = line
480 }
481 groups = append(groups, group)
482 }
483
484 return groups
485 }
486
View as plain text