imports.go

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