go.tools/astutil: a package for common AST utilities.
Lifts the import management utilities from gofix into a package, so they can be used by goimports. R=bradfitz CC=golang-dev https://golang.org/cl/22430043
This commit is contained in:
David Crawshaw
parent
eb67edd047
commit
0b5928ea2a
|
|
@ -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 "<nil>":
|
||||
// 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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue