github.com/torresashjian/cli@v0.10.1-0.20210916231452-89080fe7069c/util/ast.go (about) 1 package util 2 3 import ( 4 "fmt" 5 "go/ast" 6 "go/token" 7 "strconv" 8 "strings" 9 ) 10 11 // Copyright 2013 The Go Authors. All rights reserved. 12 // Use of this source code is governed by a BSD-style 13 // license that can be found in the LICENSE file. 14 15 // code from: https://github.com/golang/tools/blob/master/go/ast/astutil/imports.go 16 17 // AddImport adds the import path to the file f, if absent. 18 func AddImport(fset *token.FileSet, f *ast.File, ipath string) (added bool) { 19 return addNamedImport(fset, f, "_", ipath) 20 } 21 22 // AddNamedImport adds the import path to the file f, if absent. 23 // If name is not empty, it is used to rename the import. 24 // 25 // For example, calling 26 // AddNamedImport(fset, f, "pathpkg", "path") 27 // adds 28 // import pathpkg "path" 29 func addNamedImport(fset *token.FileSet, f *ast.File, name, ipath string) (added bool) { 30 if imports(f, ipath) { 31 return false 32 } 33 34 newImport := &ast.ImportSpec{ 35 Path: &ast.BasicLit{ 36 Kind: token.STRING, 37 Value: strconv.Quote(ipath), 38 }, 39 } 40 if name != "" { 41 newImport.Name = &ast.Ident{Name: name} 42 } 43 44 // Find an import decl to add to. 45 // The goal is to find an existing import 46 // whose import path has the longest shared 47 // prefix with ipath. 48 var ( 49 bestMatch = -1 // length of longest shared prefix 50 lastImport = -1 // index in f.Decls of the file's final import decl 51 impDecl *ast.GenDecl // import decl containing the best match 52 impIndex = -1 // spec index in impDecl containing the best match 53 54 isThirdPartyPath = isThirdParty(ipath) 55 ) 56 for i, decl := range f.Decls { 57 gen, ok := decl.(*ast.GenDecl) 58 if ok && gen.Tok == token.IMPORT { 59 lastImport = i 60 // Do not add to import "C", to avoid disrupting the 61 // association with its doc comment, breaking cgo. 62 if declImports(gen, "C") { 63 continue 64 } 65 66 // Match an empty import decl if that's all that is available. 67 if len(gen.Specs) == 0 && bestMatch == -1 { 68 impDecl = gen 69 } 70 71 // Compute longest shared prefix with imports in this group and find best 72 // matched import spec. 73 // 1. Always prefer import spec with longest shared prefix. 74 // 2. While match length is 0, 75 // - for stdlib package: prefer first import spec. 76 // - for third party package: prefer first third party import spec. 77 // We cannot use last import spec as best match for third party package 78 // because grouped imports are usually placed last by goimports -local 79 // flag. 80 // See issue #19190. 81 seenAnyThirdParty := false 82 for j, spec := range gen.Specs { 83 impspec := spec.(*ast.ImportSpec) 84 p := importPath(impspec) 85 n := matchLen(p, ipath) 86 if n > bestMatch || (bestMatch == 0 && !seenAnyThirdParty && isThirdPartyPath) { 87 bestMatch = n 88 impDecl = gen 89 impIndex = j 90 } 91 seenAnyThirdParty = seenAnyThirdParty || isThirdParty(p) 92 } 93 } 94 } 95 96 // If no import decl found, add one after the last import. 97 if impDecl == nil { 98 impDecl = &ast.GenDecl{ 99 Tok: token.IMPORT, 100 } 101 if lastImport >= 0 { 102 impDecl.TokPos = f.Decls[lastImport].End() 103 } else { 104 // There are no existing imports. 105 // Our new import, preceded by a blank line, goes after the package declaration 106 // and after the comment, if any, that starts on the same line as the 107 // package declaration. 108 impDecl.TokPos = f.Package 109 110 file := fset.File(f.Package) 111 pkgLine := file.Line(f.Package) 112 for _, c := range f.Comments { 113 if file.Line(c.Pos()) > pkgLine { 114 break 115 } 116 // +2 for a blank line 117 impDecl.TokPos = c.End() + 2 118 } 119 } 120 f.Decls = append(f.Decls, nil) 121 copy(f.Decls[lastImport+2:], f.Decls[lastImport+1:]) 122 f.Decls[lastImport+1] = impDecl 123 } 124 125 // Insert new import at insertAt. 126 insertAt := 0 127 if impIndex >= 0 { 128 // insert after the found import 129 insertAt = impIndex + 1 130 } 131 impDecl.Specs = append(impDecl.Specs, nil) 132 copy(impDecl.Specs[insertAt+1:], impDecl.Specs[insertAt:]) 133 impDecl.Specs[insertAt] = newImport 134 pos := impDecl.Pos() 135 if insertAt > 0 { 136 // If there is a comment after an existing import, preserve the comment 137 // position by adding the new import after the comment. 138 if spec, ok := impDecl.Specs[insertAt-1].(*ast.ImportSpec); ok && spec.Comment != nil { 139 pos = spec.Comment.End() 140 } else { 141 // Assign same position as the previous import, 142 // so that the sorter sees it as being in the same block. 143 pos = impDecl.Specs[insertAt-1].Pos() 144 } 145 } 146 if newImport.Name != nil { 147 newImport.Name.NamePos = pos 148 } 149 newImport.Path.ValuePos = pos 150 newImport.EndPos = pos 151 152 // Clean up parens. impDecl contains at least one spec. 153 if len(impDecl.Specs) == 1 { 154 // Remove unneeded parens. 155 impDecl.Lparen = token.NoPos 156 } else if !impDecl.Lparen.IsValid() { 157 // impDecl needs parens added. 158 impDecl.Lparen = impDecl.Specs[0].Pos() 159 } 160 161 f.Imports = append(f.Imports, newImport) 162 163 if len(f.Decls) <= 1 { 164 return true 165 } 166 167 // Merge all the import declarations into the first one. 168 var first *ast.GenDecl 169 for i := 0; i < len(f.Decls); i++ { 170 decl := f.Decls[i] 171 gen, ok := decl.(*ast.GenDecl) 172 if !ok || gen.Tok != token.IMPORT || declImports(gen, "C") { 173 continue 174 } 175 if first == nil { 176 first = gen 177 continue // Don't touch the first one. 178 } 179 // We now know there is more than one package in this import 180 // declaration. Ensure that it ends up parenthesized. 181 first.Lparen = first.Pos() 182 // Move the imports of the other import declaration to the first one. 183 for _, spec := range gen.Specs { 184 spec.(*ast.ImportSpec).Path.ValuePos = first.Pos() 185 first.Specs = append(first.Specs, spec) 186 } 187 f.Decls = append(f.Decls[:i], f.Decls[i+1:]...) 188 i-- 189 } 190 191 return true 192 } 193 194 func isThirdParty(importPath string) bool { 195 // Third party package import path usually contains "." (".com", ".org", ...) 196 // This logic is taken from golang.org/x/tools/imports package. 197 return strings.Contains(importPath, ".") 198 } 199 200 // DeleteImport deletes the import path from the file f, if present. 201 func DeleteImport(fset *token.FileSet, f *ast.File, path string) (deleted bool) { 202 return deleteNamedImport(fset, f, "_", path) 203 } 204 205 // DeleteNamedImport deletes the import with the given name and path from the file f, if present. 206 func deleteNamedImport(fset *token.FileSet, f *ast.File, name, path string) (deleted bool) { 207 var delspecs []*ast.ImportSpec 208 var delcomments []*ast.CommentGroup 209 210 // Find the import nodes that import path, if any. 211 for i := 0; i < len(f.Decls); i++ { 212 decl := f.Decls[i] 213 gen, ok := decl.(*ast.GenDecl) 214 if !ok || gen.Tok != token.IMPORT { 215 continue 216 } 217 for j := 0; j < len(gen.Specs); j++ { 218 spec := gen.Specs[j] 219 impspec := spec.(*ast.ImportSpec) 220 if impspec.Name == nil && name != "" { 221 continue 222 } 223 if impspec.Name != nil && impspec.Name.Name != name { 224 continue 225 } 226 if importPath(impspec) != path { 227 continue 228 } 229 230 // We found an import spec that imports path. 231 // Delete it. 232 delspecs = append(delspecs, impspec) 233 deleted = true 234 copy(gen.Specs[j:], gen.Specs[j+1:]) 235 gen.Specs = gen.Specs[:len(gen.Specs)-1] 236 237 // If this was the last import spec in this decl, 238 // delete the decl, too. 239 if len(gen.Specs) == 0 { 240 copy(f.Decls[i:], f.Decls[i+1:]) 241 f.Decls = f.Decls[:len(f.Decls)-1] 242 i-- 243 break 244 } else if len(gen.Specs) == 1 { 245 if impspec.Doc != nil { 246 delcomments = append(delcomments, impspec.Doc) 247 } 248 if impspec.Comment != nil { 249 delcomments = append(delcomments, impspec.Comment) 250 } 251 for _, cg := range f.Comments { 252 // Found comment on the same line as the import spec. 253 if cg.End() < impspec.Pos() && fset.Position(cg.End()).Line == fset.Position(impspec.Pos()).Line { 254 delcomments = append(delcomments, cg) 255 break 256 } 257 } 258 259 spec := gen.Specs[0].(*ast.ImportSpec) 260 261 // Move the documentation right after the import decl. 262 if spec.Doc != nil { 263 for fset.Position(gen.TokPos).Line+1 < fset.Position(spec.Doc.Pos()).Line { 264 fset.File(gen.TokPos).MergeLine(fset.Position(gen.TokPos).Line) 265 } 266 } 267 for _, cg := range f.Comments { 268 if cg.End() < spec.Pos() && fset.Position(cg.End()).Line == fset.Position(spec.Pos()).Line { 269 for fset.Position(gen.TokPos).Line+1 < fset.Position(spec.Pos()).Line { 270 fset.File(gen.TokPos).MergeLine(fset.Position(gen.TokPos).Line) 271 } 272 break 273 } 274 } 275 } 276 if j > 0 { 277 lastImpspec := gen.Specs[j-1].(*ast.ImportSpec) 278 lastLine := fset.Position(lastImpspec.Path.ValuePos).Line 279 line := fset.Position(impspec.Path.ValuePos).Line 280 281 // We deleted an entry but now there may be 282 // a blank line-sized hole where the import was. 283 if line-lastLine > 1 { 284 // There was a blank line immediately preceding the deleted import, 285 // so there's no need to close the hole. 286 // Do nothing. 287 } else if line != fset.File(gen.Rparen).LineCount() { 288 // There was no blank line. Close the hole. 289 fset.File(gen.Rparen).MergeLine(line) 290 } 291 } 292 j-- 293 } 294 } 295 296 // Delete imports from f.Imports. 297 for i := 0; i < len(f.Imports); i++ { 298 imp := f.Imports[i] 299 for j, del := range delspecs { 300 if imp == del { 301 copy(f.Imports[i:], f.Imports[i+1:]) 302 f.Imports = f.Imports[:len(f.Imports)-1] 303 copy(delspecs[j:], delspecs[j+1:]) 304 delspecs = delspecs[:len(delspecs)-1] 305 i-- 306 break 307 } 308 } 309 } 310 311 // Delete comments from f.Comments. 312 for i := 0; i < len(f.Comments); i++ { 313 cg := f.Comments[i] 314 for j, del := range delcomments { 315 if cg == del { 316 copy(f.Comments[i:], f.Comments[i+1:]) 317 f.Comments = f.Comments[:len(f.Comments)-1] 318 copy(delcomments[j:], delcomments[j+1:]) 319 delcomments = delcomments[:len(delcomments)-1] 320 i-- 321 break 322 } 323 } 324 } 325 326 if len(delspecs) > 0 { 327 panic(fmt.Sprintf("deleted specs from Decls but not Imports: %v", delspecs)) 328 } 329 330 return 331 } 332 333 // 334 //// RewriteImport rewrites any import of path oldPath to path newPath. 335 //func RewriteImport(fset *token.FileSet, f *ast.File, oldPath, newPath string) (rewrote bool) { 336 // for _, imp := range f.Imports { 337 // if importPath(imp) == oldPath { 338 // rewrote = true 339 // // record old End, because the default is to compute 340 // // it using the length of imp.Path.Value. 341 // imp.EndPos = imp.End() 342 // imp.Path.Value = strconv.Quote(newPath) 343 // } 344 // } 345 // return 346 //} 347 // 348 //// UsesImport reports whether a given import is used. 349 //func UsesImport(f *ast.File, path string) (used bool) { 350 // spec := importSpec(f, path) 351 // if spec == nil { 352 // return 353 // } 354 // 355 // name := spec.Name.String() 356 // switch name { 357 // case "<nil>": 358 // // If the package name is not explicitly specified, 359 // // make an educated guess. This is not guaranteed to be correct. 360 // lastSlash := strings.LastIndex(path, "/") 361 // if lastSlash == -1 { 362 // name = path 363 // } else { 364 // name = path[lastSlash+1:] 365 // } 366 // case "_", ".": 367 // // Not sure if this import is used - err on the side of caution. 368 // return true 369 // } 370 // 371 // ast.Walk(visitFn(func(n ast.Node) { 372 // sel, ok := n.(*ast.SelectorExpr) 373 // if ok && isTopName(sel.X, name) { 374 // used = true 375 // } 376 // }), f) 377 // 378 // return 379 //} 380 381 //type visitFn func(node ast.Node) 382 // 383 //func (fn visitFn) Visit(node ast.Node) ast.Visitor { 384 // fn(node) 385 // return fn 386 //} 387 388 // imports returns true if f imports path. 389 func imports(f *ast.File, path string) bool { 390 return importSpec(f, path) != nil 391 } 392 393 // importSpec returns the import spec if f imports path, 394 // or nil otherwise. 395 func importSpec(f *ast.File, path string) *ast.ImportSpec { 396 for _, s := range f.Imports { 397 if importPath(s) == path { 398 return s 399 } 400 } 401 return nil 402 } 403 404 // importPath returns the unquoted import path of s, 405 // or "" if the path is not properly quoted. 406 func importPath(s *ast.ImportSpec) string { 407 t, err := strconv.Unquote(s.Path.Value) 408 if err == nil { 409 return t 410 } 411 return "" 412 } 413 414 // declImports reports whether gen contains an import of path. 415 func declImports(gen *ast.GenDecl, path string) bool { 416 if gen.Tok != token.IMPORT { 417 return false 418 } 419 for _, spec := range gen.Specs { 420 impspec := spec.(*ast.ImportSpec) 421 if importPath(impspec) == path { 422 return true 423 } 424 } 425 return false 426 } 427 428 // matchLen returns the length of the longest path segment prefix shared by x and y. 429 func matchLen(x, y string) int { 430 n := 0 431 for i := 0; i < len(x) && i < len(y) && x[i] == y[i]; i++ { 432 if x[i] == '/' { 433 n++ 434 } 435 } 436 return n 437 } 438 439 //// isTopName returns true if n is a top-level unresolved identifier with the given name. 440 //func isTopName(n ast.Expr, name string) bool { 441 // id, ok := n.(*ast.Ident) 442 // return ok && id.Name == name && id.Obj == nil 443 //} 444 // 445 //// Imports returns the file imports grouped by paragraph. 446 //func Imports(fset *token.FileSet, f *ast.File) [][]*ast.ImportSpec { 447 // var groups [][]*ast.ImportSpec 448 // 449 // for _, decl := range f.Decls { 450 // genDecl, ok := decl.(*ast.GenDecl) 451 // if !ok || genDecl.Tok != token.IMPORT { 452 // break 453 // } 454 // 455 // group := []*ast.ImportSpec{} 456 // 457 // var lastLine int 458 // for _, spec := range genDecl.Specs { 459 // importSpec := spec.(*ast.ImportSpec) 460 // pos := importSpec.Path.ValuePos 461 // line := fset.Position(pos).Line 462 // if lastLine > 0 && pos > 0 && line-lastLine > 1 { 463 // groups = append(groups, group) 464 // group = []*ast.ImportSpec{} 465 // } 466 // group = append(group, importSpec) 467 // lastLine = line 468 // } 469 // groups = append(groups, group) 470 // } 471 // 472 // return groups 473 //}