diff --git a/astutil/imports.go b/astutil/imports.go new file mode 100644 index 00000000..2fa7a518 --- /dev/null +++ b/astutil/imports.go @@ -0,0 +1,327 @@ +// Package astutil contains common utilities for working with the Go AST. +package astutil + +import ( + "go/ast" + "go/token" + "path" + "strconv" + "strings" +) + +// AddImport adds the import path to the file f, if absent. +func AddImport(f *ast.File, ipath string) (added bool) { + if imports(f, ipath) { + return false + } + + // Determine name of import. + // Assume added imports follow convention of using last element. + _, name := path.Split(ipath) + + // Rename any conflicting top-level references from name to name_. + renameTop(f, name, name+"_") + + newImport := &ast.ImportSpec{ + Path: &ast.BasicLit{ + Kind: token.STRING, + Value: strconv.Quote(ipath), + }, + } + + // Find an import decl to add to. + var ( + bestMatch = -1 + lastImport = -1 + impDecl *ast.GenDecl + impIndex = -1 + ) + for i, decl := range f.Decls { + gen, ok := decl.(*ast.GenDecl) + if ok && gen.Tok == token.IMPORT { + lastImport = i + // Do not add to import "C", to avoid disrupting the + // association with its doc comment, breaking cgo. + if declImports(gen, "C") { + continue + } + + // Compute longest shared prefix with imports in this block. + for j, spec := range gen.Specs { + impspec := spec.(*ast.ImportSpec) + n := matchLen(importPath(impspec), ipath) + if n > bestMatch { + bestMatch = n + impDecl = gen + impIndex = j + } + } + } + } + + // If no import decl found, add one after the last import. + if impDecl == nil { + impDecl = &ast.GenDecl{ + Tok: token.IMPORT, + } + f.Decls = append(f.Decls, nil) + copy(f.Decls[lastImport+2:], f.Decls[lastImport+1:]) + f.Decls[lastImport+1] = impDecl + } + + // Ensure the import decl has parentheses, if needed. + if len(impDecl.Specs) > 0 && !impDecl.Lparen.IsValid() { + impDecl.Lparen = impDecl.Pos() + } + + insertAt := impIndex + 1 + if insertAt == 0 { + insertAt = len(impDecl.Specs) + } + impDecl.Specs = append(impDecl.Specs, nil) + copy(impDecl.Specs[insertAt+1:], impDecl.Specs[insertAt:]) + impDecl.Specs[insertAt] = newImport + if insertAt > 0 { + // Assign same position as the previous import, + // so that the sorter sees it as being in the same block. + prev := impDecl.Specs[insertAt-1] + newImport.Path.ValuePos = prev.Pos() + newImport.EndPos = prev.Pos() + } + + f.Imports = append(f.Imports, newImport) + return true +} + +// DeleteImport deletes the import path from the file f, if present. +func DeleteImport(f *ast.File, path string) (deleted bool) { + oldImport := importSpec(f, path) + + // Find the import node that imports path, if any. + for i, decl := range f.Decls { + gen, ok := decl.(*ast.GenDecl) + if !ok || gen.Tok != token.IMPORT { + continue + } + for j, spec := range gen.Specs { + impspec := spec.(*ast.ImportSpec) + if oldImport != impspec { + continue + } + + // We found an import spec that imports path. + // Delete it. + deleted = true + copy(gen.Specs[j:], gen.Specs[j+1:]) + gen.Specs = gen.Specs[:len(gen.Specs)-1] + + // If this was the last import spec in this decl, + // delete the decl, too. + if len(gen.Specs) == 0 { + copy(f.Decls[i:], f.Decls[i+1:]) + f.Decls = f.Decls[:len(f.Decls)-1] + } else if len(gen.Specs) == 1 { + gen.Lparen = token.NoPos // drop parens + } + if j > 0 { + // We deleted an entry but now there will be + // a blank line-sized hole where the import was. + // Close the hole by making the previous + // import appear to "end" where this one did. + gen.Specs[j-1].(*ast.ImportSpec).EndPos = impspec.End() + } + break + } + } + + // Delete it from f.Imports. + for i, imp := range f.Imports { + if imp == oldImport { + copy(f.Imports[i:], f.Imports[i+1:]) + f.Imports = f.Imports[:len(f.Imports)-1] + break + } + } + + return +} + +// RewriteImport rewrites any import of path oldPath to path newPath. +func RewriteImport(f *ast.File, oldPath, newPath string) (rewrote bool) { + for _, imp := range f.Imports { + if importPath(imp) == oldPath { + rewrote = true + // record old End, because the default is to compute + // it using the length of imp.Path.Value. + imp.EndPos = imp.End() + imp.Path.Value = strconv.Quote(newPath) + } + } + return +} + +// UsesImport reports whether a given import is used. +func UsesImport(f *ast.File, path string) (used bool) { + spec := importSpec(f, path) + if spec == nil { + return + } + + name := spec.Name.String() + switch name { + case "": + // If the package name is not explicitly specified, + // make an educated guess. This is not guaranteed to be correct. + lastSlash := strings.LastIndex(path, "/") + if lastSlash == -1 { + name = path + } else { + name = path[lastSlash+1:] + } + case "_", ".": + // Not sure if this import is used - err on the side of caution. + return true + } + + ast.Walk(visitFn(func(n ast.Node) { + sel, ok := n.(*ast.SelectorExpr) + if ok && isTopName(sel.X, name) { + used = true + } + }), f) + + return +} + +type visitFn func(node ast.Node) + +func (fn visitFn) Visit(node ast.Node) ast.Visitor { + fn(node) + return fn +} + +// imports returns true if f imports path. +func imports(f *ast.File, path string) bool { + return importSpec(f, path) != nil +} + +// importSpec returns the import spec if f imports path, +// or nil otherwise. +func importSpec(f *ast.File, path string) *ast.ImportSpec { + for _, s := range f.Imports { + if importPath(s) == path { + return s + } + } + return nil +} + +// importPath returns the unquoted import path of s, +// or "" if the path is not properly quoted. +func importPath(s *ast.ImportSpec) string { + t, err := strconv.Unquote(s.Path.Value) + if err == nil { + return t + } + return "" +} + +// declImports reports whether gen contains an import of path. +func declImports(gen *ast.GenDecl, path string) bool { + if gen.Tok != token.IMPORT { + return false + } + for _, spec := range gen.Specs { + impspec := spec.(*ast.ImportSpec) + if importPath(impspec) == path { + return true + } + } + return false +} + +// renameTop renames all references to the top-level name old. +// It returns true if it makes any changes. +func renameTop(f *ast.File, old, new string) bool { + var fixed bool + + // Rename any conflicting imports + // (assuming package name is last element of path). + for _, s := range f.Imports { + if s.Name != nil { + if s.Name.Name == old { + s.Name.Name = new + fixed = true + } + } else { + _, thisName := path.Split(importPath(s)) + if thisName == old { + s.Name = ast.NewIdent(new) + fixed = true + } + } + } + + // Rename any top-level declarations. + for _, d := range f.Decls { + switch d := d.(type) { + case *ast.FuncDecl: + if d.Recv == nil && d.Name.Name == old { + d.Name.Name = new + d.Name.Obj.Name = new + fixed = true + } + case *ast.GenDecl: + for _, s := range d.Specs { + switch s := s.(type) { + case *ast.TypeSpec: + if s.Name.Name == old { + s.Name.Name = new + s.Name.Obj.Name = new + fixed = true + } + case *ast.ValueSpec: + for _, n := range s.Names { + if n.Name == old { + n.Name = new + n.Obj.Name = new + fixed = true + } + } + } + } + } + } + + // Rename top-level old to new, both unresolved names + // (probably defined in another file) and names that resolve + // to a declaration we renamed. + ast.Walk(visitFn(func(n ast.Node) { + id, ok := n.(*ast.Ident) + if ok && isTopName(id, old) { + id.Name = new + fixed = true + } + if ok && id.Obj != nil && id.Name == old && id.Obj.Name == new { + id.Name = id.Obj.Name + fixed = true + } + }), f) + + return fixed +} + +// matchLen returns the length of the longest prefix shared by x and y. +func matchLen(x, y string) int { + i := 0 + for i < len(x) && i < len(y) && x[i] == y[i] { + i++ + } + return i +} + +// isTopName returns true if n is a top-level unresolved identifier with the given name. +func isTopName(n ast.Expr, name string) bool { + id, ok := n.(*ast.Ident) + return ok && id.Name == name && id.Obj == nil +} diff --git a/astutil/imports_test.go b/astutil/imports_test.go new file mode 100644 index 00000000..86d6775d --- /dev/null +++ b/astutil/imports_test.go @@ -0,0 +1,459 @@ +package astutil + +import ( + "bytes" + "go/ast" + "go/format" + "go/parser" + "go/token" + + "testing" +) + +var fset = token.NewFileSet() + +func parse(t *testing.T, name, in string) *ast.File { + file, err := parser.ParseFile(fset, name, in, parser.ParseComments) + if err != nil { + t.Fatalf("%s parse: %v", name, err) + } + return file +} + +func print(t *testing.T, name string, f *ast.File) string { + var buf bytes.Buffer + if err := format.Node(&buf, fset, f); err != nil { + t.Fatalf("%s gofmt: %v", name, err) + } + return string(buf.Bytes()) +} + +type test struct { + name string + pkg string + in string + out string +} + +var addTests = []test{ + { + name: "leave os alone", + pkg: "os", + in: `package main + +import ( + "os" +) +`, + out: `package main + +import ( + "os" +) +`, + }, + { + name: "import.1", + pkg: "os", + in: `package main +`, + out: `package main + +import "os" +`, + }, + { + name: "import.2", + pkg: "os", + in: `package main + +// Comment +import "C" +`, + out: `package main + +// Comment +import "C" +import "os" +`, + }, + { + name: "import.3", + pkg: "os", + in: `package main + +// Comment +import "C" + +import ( + "io" + "utf8" +) +`, + out: `package main + +// Comment +import "C" + +import ( + "io" + "os" + "utf8" +) +`, + }, + { + name: "import.17", + pkg: "x/y/z", + in: `package main + +// Comment +import "C" + +import ( + "a" + "b" + + "x/w" + + "d/f" +) +`, + out: `package main + +// Comment +import "C" + +import ( + "a" + "b" + + "x/w" + "x/y/z" + + "d/f" +) +`, + }, +} + +func TestAddImport(t *testing.T) { + for _, test := range addTests { + file := parse(t, test.name, test.in) + AddImport(file, test.pkg) + if got := print(t, test.name, file); got != test.out { + t.Errorf("%s:\ngot: %s\nwant: %s", test.name, got, test.out) + } + } +} + +var deleteTests = []test{ + { + name: "import.4", + pkg: "os", + in: `package main + +import ( + "os" +) +`, + out: `package main +`, + }, + { + name: "import.5", + pkg: "os", + in: `package main + +// Comment +import "C" +import "os" +`, + out: `package main + +// Comment +import "C" +`, + }, + { + name: "import.6", + pkg: "os", + in: `package main + +// Comment +import "C" + +import ( + "io" + "os" + "utf8" +) +`, + out: `package main + +// Comment +import "C" + +import ( + "io" + "utf8" +) +`, + }, + { + name: "import.7", + pkg: "io", + in: `package main + +import ( + "io" // a + "os" // b + "utf8" // c +) +`, + out: `package main + +import ( + // a + "os" // b + "utf8" // c +) +`, + }, + { + name: "import.8", + pkg: "os", + in: `package main + +import ( + "io" // a + "os" // b + "utf8" // c +) +`, + out: `package main + +import ( + "io" // a + // b + "utf8" // c +) +`, + }, + { + name: "import.9", + pkg: "utf8", + in: `package main + +import ( + "io" // a + "os" // b + "utf8" // c +) +`, + out: `package main + +import ( + "io" // a + "os" // b + // c +) +`, + }, + { + name: "import.10", + pkg: "io", + in: `package main + +import ( + "io" + "os" + "utf8" +) +`, + out: `package main + +import ( + "os" + "utf8" +) +`, + }, + { + name: "import.11", + pkg: "os", + in: `package main + +import ( + "io" + "os" + "utf8" +) +`, + out: `package main + +import ( + "io" + "utf8" +) +`, + }, + { + name: "import.12", + pkg: "utf8", + in: `package main + +import ( + "io" + "os" + "utf8" +) +`, + out: `package main + +import ( + "io" + "os" +) +`, + }, +} + +func TestDeleteImport(t *testing.T) { + for _, test := range deleteTests { + file := parse(t, test.name, test.in) + DeleteImport(file, test.pkg) + if got := print(t, test.name, file); got != test.out { + t.Errorf("%s:\ngot: %s\nwant: %s", test.name, got, test.out) + } + } +} + +type rewriteTest struct { + name string + srcPkg string + dstPkg string + in string + out string +} + +var rewriteTests = []rewriteTest{ + { + name: "import.13", + srcPkg: "utf8", + dstPkg: "encoding/utf8", + in: `package main + +import ( + "io" + "os" + "utf8" // thanks ken +) +`, + out: `package main + +import ( + "encoding/utf8" // thanks ken + "io" + "os" +) +`, + }, + { + name: "import.14", + srcPkg: "asn1", + dstPkg: "encoding/asn1", + in: `package main + +import ( + "asn1" + "crypto" + "crypto/rsa" + _ "crypto/sha1" + "crypto/x509" + "crypto/x509/pkix" + "time" +) + +var x = 1 +`, + out: `package main + +import ( + "crypto" + "crypto/rsa" + _ "crypto/sha1" + "crypto/x509" + "crypto/x509/pkix" + "encoding/asn1" + "time" +) + +var x = 1 +`, + }, + { + name: "import.15", + srcPkg: "url", + dstPkg: "net/url", + in: `package main + +import ( + "bufio" + "net" + "path" + "url" +) + +var x = 1 // comment on x, not on url +`, + out: `package main + +import ( + "bufio" + "net" + "net/url" + "path" +) + +var x = 1 // comment on x, not on url +`, + }, + { + name: "import.16", + srcPkg: "http", + dstPkg: "net/http", + in: `package main + +import ( + "flag" + "http" + "log" + "text/template" +) + +var addr = flag.String("addr", ":1718", "http service address") // Q=17, R=18 +`, + out: `package main + +import ( + "flag" + "log" + "net/http" + "text/template" +) + +var addr = flag.String("addr", ":1718", "http service address") // Q=17, R=18 +`, + }, +} + +func TestRewriteImport(t *testing.T) { + for _, test := range rewriteTests { + file := parse(t, test.name, test.in) + RewriteImport(file, test.srcPkg, test.dstPkg) + if got := print(t, test.name, file); got != test.out { + t.Errorf("%s:\ngot: %s\nwant: %s", test.name, got, test.out) + } + } +}