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