You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

imports.go 11KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408
  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. //go:generate go run mkstdlib.go
  5. // Package imports implements a Go pretty-printer (like package "go/format")
  6. // that also adds or removes import statements as necessary.
  7. package imports
  8. import (
  9. "bufio"
  10. "bytes"
  11. "context"
  12. "fmt"
  13. "go/ast"
  14. "go/build"
  15. "go/format"
  16. "go/parser"
  17. "go/printer"
  18. "go/token"
  19. "io"
  20. "io/ioutil"
  21. "os"
  22. "regexp"
  23. "strconv"
  24. "strings"
  25. "golang.org/x/tools/go/ast/astutil"
  26. )
  27. // Options is golang.org/x/tools/imports.Options with extra internal-only options.
  28. type Options struct {
  29. Env *ProcessEnv // The environment to use. Note: this contains the cached module and filesystem state.
  30. Fragment bool // Accept fragment of a source file (no package statement)
  31. AllErrors bool // Report all errors (not just the first 10 on different lines)
  32. Comments bool // Print comments (true if nil *Options provided)
  33. TabIndent bool // Use tabs for indent (true if nil *Options provided)
  34. TabWidth int // Tab width (8 if nil *Options provided)
  35. FormatOnly bool // Disable the insertion and deletion of imports
  36. }
  37. // Process implements golang.org/x/tools/imports.Process with explicit context in env.
  38. func Process(filename string, src []byte, opt *Options) (formatted []byte, err error) {
  39. src, opt, err = initialize(filename, src, opt)
  40. if err != nil {
  41. return nil, err
  42. }
  43. fileSet := token.NewFileSet()
  44. file, adjust, err := parse(fileSet, filename, src, opt)
  45. if err != nil {
  46. return nil, err
  47. }
  48. if !opt.FormatOnly {
  49. if err := fixImports(fileSet, file, filename, opt.Env); err != nil {
  50. return nil, err
  51. }
  52. }
  53. return formatFile(fileSet, file, src, adjust, opt)
  54. }
  55. // FixImports returns a list of fixes to the imports that, when applied,
  56. // will leave the imports in the same state as Process.
  57. //
  58. // Note that filename's directory influences which imports can be chosen,
  59. // so it is important that filename be accurate.
  60. func FixImports(filename string, src []byte, opt *Options) (fixes []*ImportFix, err error) {
  61. src, opt, err = initialize(filename, src, opt)
  62. if err != nil {
  63. return nil, err
  64. }
  65. fileSet := token.NewFileSet()
  66. file, _, err := parse(fileSet, filename, src, opt)
  67. if err != nil {
  68. return nil, err
  69. }
  70. return getFixes(fileSet, file, filename, opt.Env)
  71. }
  72. // ApplyFixes applies all of the fixes to the file and formats it. extraMode
  73. // is added in when parsing the file.
  74. func ApplyFixes(fixes []*ImportFix, filename string, src []byte, opt *Options, extraMode parser.Mode) (formatted []byte, err error) {
  75. src, opt, err = initialize(filename, src, opt)
  76. if err != nil {
  77. return nil, err
  78. }
  79. // Don't use parse() -- we don't care about fragments or statement lists
  80. // here, and we need to work with unparseable files.
  81. fileSet := token.NewFileSet()
  82. parserMode := parser.Mode(0)
  83. if opt.Comments {
  84. parserMode |= parser.ParseComments
  85. }
  86. if opt.AllErrors {
  87. parserMode |= parser.AllErrors
  88. }
  89. parserMode |= extraMode
  90. file, err := parser.ParseFile(fileSet, filename, src, parserMode)
  91. if file == nil {
  92. return nil, err
  93. }
  94. // Apply the fixes to the file.
  95. apply(fileSet, file, fixes)
  96. return formatFile(fileSet, file, src, nil, opt)
  97. }
  98. // GetAllCandidates gets all of the packages starting with prefix that can be
  99. // imported by filename, sorted by import path.
  100. func GetAllCandidates(ctx context.Context, callback func(ImportFix), searchPrefix, filename, filePkg string, opt *Options) error {
  101. _, opt, err := initialize(filename, []byte{}, opt)
  102. if err != nil {
  103. return err
  104. }
  105. return getAllCandidates(ctx, callback, searchPrefix, filename, filePkg, opt.Env)
  106. }
  107. // GetPackageExports returns all known packages with name pkg and their exports.
  108. func GetPackageExports(ctx context.Context, callback func(PackageExport), searchPkg, filename, filePkg string, opt *Options) error {
  109. _, opt, err := initialize(filename, []byte{}, opt)
  110. if err != nil {
  111. return err
  112. }
  113. return getPackageExports(ctx, callback, searchPkg, filename, filePkg, opt.Env)
  114. }
  115. // initialize sets the values for opt and src.
  116. // If they are provided, they are not changed. Otherwise opt is set to the
  117. // default values and src is read from the file system.
  118. func initialize(filename string, src []byte, opt *Options) ([]byte, *Options, error) {
  119. // Use defaults if opt is nil.
  120. if opt == nil {
  121. opt = &Options{Comments: true, TabIndent: true, TabWidth: 8}
  122. }
  123. // Set the env if the user has not provided it.
  124. if opt.Env == nil {
  125. opt.Env = &ProcessEnv{
  126. GOPATH: build.Default.GOPATH,
  127. GOROOT: build.Default.GOROOT,
  128. GOFLAGS: os.Getenv("GOFLAGS"),
  129. GO111MODULE: os.Getenv("GO111MODULE"),
  130. GOPROXY: os.Getenv("GOPROXY"),
  131. GOSUMDB: os.Getenv("GOSUMDB"),
  132. }
  133. }
  134. if src == nil {
  135. b, err := ioutil.ReadFile(filename)
  136. if err != nil {
  137. return nil, nil, err
  138. }
  139. src = b
  140. }
  141. return src, opt, nil
  142. }
  143. func formatFile(fileSet *token.FileSet, file *ast.File, src []byte, adjust func(orig []byte, src []byte) []byte, opt *Options) ([]byte, error) {
  144. mergeImports(opt.Env, fileSet, file)
  145. sortImports(opt.Env, fileSet, file)
  146. imps := astutil.Imports(fileSet, file)
  147. var spacesBefore []string // import paths we need spaces before
  148. for _, impSection := range imps {
  149. // Within each block of contiguous imports, see if any
  150. // import lines are in different group numbers. If so,
  151. // we'll need to put a space between them so it's
  152. // compatible with gofmt.
  153. lastGroup := -1
  154. for _, importSpec := range impSection {
  155. importPath, _ := strconv.Unquote(importSpec.Path.Value)
  156. groupNum := importGroup(opt.Env, importPath)
  157. if groupNum != lastGroup && lastGroup != -1 {
  158. spacesBefore = append(spacesBefore, importPath)
  159. }
  160. lastGroup = groupNum
  161. }
  162. }
  163. printerMode := printer.UseSpaces
  164. if opt.TabIndent {
  165. printerMode |= printer.TabIndent
  166. }
  167. printConfig := &printer.Config{Mode: printerMode, Tabwidth: opt.TabWidth}
  168. var buf bytes.Buffer
  169. err := printConfig.Fprint(&buf, fileSet, file)
  170. if err != nil {
  171. return nil, err
  172. }
  173. out := buf.Bytes()
  174. if adjust != nil {
  175. out = adjust(src, out)
  176. }
  177. if len(spacesBefore) > 0 {
  178. out, err = addImportSpaces(bytes.NewReader(out), spacesBefore)
  179. if err != nil {
  180. return nil, err
  181. }
  182. }
  183. out, err = format.Source(out)
  184. if err != nil {
  185. return nil, err
  186. }
  187. return out, nil
  188. }
  189. // parse parses src, which was read from filename,
  190. // as a Go source file or statement list.
  191. func parse(fset *token.FileSet, filename string, src []byte, opt *Options) (*ast.File, func(orig, src []byte) []byte, error) {
  192. parserMode := parser.Mode(0)
  193. if opt.Comments {
  194. parserMode |= parser.ParseComments
  195. }
  196. if opt.AllErrors {
  197. parserMode |= parser.AllErrors
  198. }
  199. // Try as whole source file.
  200. file, err := parser.ParseFile(fset, filename, src, parserMode)
  201. if err == nil {
  202. return file, nil, nil
  203. }
  204. // If the error is that the source file didn't begin with a
  205. // package line and we accept fragmented input, fall through to
  206. // try as a source fragment. Stop and return on any other error.
  207. if !opt.Fragment || !strings.Contains(err.Error(), "expected 'package'") {
  208. return nil, nil, err
  209. }
  210. // If this is a declaration list, make it a source file
  211. // by inserting a package clause.
  212. // Insert using a ;, not a newline, so that parse errors are on
  213. // the correct line.
  214. const prefix = "package main;"
  215. psrc := append([]byte(prefix), src...)
  216. file, err = parser.ParseFile(fset, filename, psrc, parserMode)
  217. if err == nil {
  218. // Gofmt will turn the ; into a \n.
  219. // Do that ourselves now and update the file contents,
  220. // so that positions and line numbers are correct going forward.
  221. psrc[len(prefix)-1] = '\n'
  222. fset.File(file.Package).SetLinesForContent(psrc)
  223. // If a main function exists, we will assume this is a main
  224. // package and leave the file.
  225. if containsMainFunc(file) {
  226. return file, nil, nil
  227. }
  228. adjust := func(orig, src []byte) []byte {
  229. // Remove the package clause.
  230. src = src[len(prefix):]
  231. return matchSpace(orig, src)
  232. }
  233. return file, adjust, nil
  234. }
  235. // If the error is that the source file didn't begin with a
  236. // declaration, fall through to try as a statement list.
  237. // Stop and return on any other error.
  238. if !strings.Contains(err.Error(), "expected declaration") {
  239. return nil, nil, err
  240. }
  241. // If this is a statement list, make it a source file
  242. // by inserting a package clause and turning the list
  243. // into a function body. This handles expressions too.
  244. // Insert using a ;, not a newline, so that the line numbers
  245. // in fsrc match the ones in src.
  246. fsrc := append(append([]byte("package p; func _() {"), src...), '}')
  247. file, err = parser.ParseFile(fset, filename, fsrc, parserMode)
  248. if err == nil {
  249. adjust := func(orig, src []byte) []byte {
  250. // Remove the wrapping.
  251. // Gofmt has turned the ; into a \n\n.
  252. src = src[len("package p\n\nfunc _() {"):]
  253. src = src[:len(src)-len("}\n")]
  254. // Gofmt has also indented the function body one level.
  255. // Remove that indent.
  256. src = bytes.Replace(src, []byte("\n\t"), []byte("\n"), -1)
  257. return matchSpace(orig, src)
  258. }
  259. return file, adjust, nil
  260. }
  261. // Failed, and out of options.
  262. return nil, nil, err
  263. }
  264. // containsMainFunc checks if a file contains a function declaration with the
  265. // function signature 'func main()'
  266. func containsMainFunc(file *ast.File) bool {
  267. for _, decl := range file.Decls {
  268. if f, ok := decl.(*ast.FuncDecl); ok {
  269. if f.Name.Name != "main" {
  270. continue
  271. }
  272. if len(f.Type.Params.List) != 0 {
  273. continue
  274. }
  275. if f.Type.Results != nil && len(f.Type.Results.List) != 0 {
  276. continue
  277. }
  278. return true
  279. }
  280. }
  281. return false
  282. }
  283. func cutSpace(b []byte) (before, middle, after []byte) {
  284. i := 0
  285. for i < len(b) && (b[i] == ' ' || b[i] == '\t' || b[i] == '\n') {
  286. i++
  287. }
  288. j := len(b)
  289. for j > 0 && (b[j-1] == ' ' || b[j-1] == '\t' || b[j-1] == '\n') {
  290. j--
  291. }
  292. if i <= j {
  293. return b[:i], b[i:j], b[j:]
  294. }
  295. return nil, nil, b[j:]
  296. }
  297. // matchSpace reformats src to use the same space context as orig.
  298. // 1) If orig begins with blank lines, matchSpace inserts them at the beginning of src.
  299. // 2) matchSpace copies the indentation of the first non-blank line in orig
  300. // to every non-blank line in src.
  301. // 3) matchSpace copies the trailing space from orig and uses it in place
  302. // of src's trailing space.
  303. func matchSpace(orig []byte, src []byte) []byte {
  304. before, _, after := cutSpace(orig)
  305. i := bytes.LastIndex(before, []byte{'\n'})
  306. before, indent := before[:i+1], before[i+1:]
  307. _, src, _ = cutSpace(src)
  308. var b bytes.Buffer
  309. b.Write(before)
  310. for len(src) > 0 {
  311. line := src
  312. if i := bytes.IndexByte(line, '\n'); i >= 0 {
  313. line, src = line[:i+1], line[i+1:]
  314. } else {
  315. src = nil
  316. }
  317. if len(line) > 0 && line[0] != '\n' { // not blank
  318. b.Write(indent)
  319. }
  320. b.Write(line)
  321. }
  322. b.Write(after)
  323. return b.Bytes()
  324. }
  325. var impLine = regexp.MustCompile(`^\s+(?:[\w\.]+\s+)?"(.+)"`)
  326. func addImportSpaces(r io.Reader, breaks []string) ([]byte, error) {
  327. var out bytes.Buffer
  328. in := bufio.NewReader(r)
  329. inImports := false
  330. done := false
  331. for {
  332. s, err := in.ReadString('\n')
  333. if err == io.EOF {
  334. break
  335. } else if err != nil {
  336. return nil, err
  337. }
  338. if !inImports && !done && strings.HasPrefix(s, "import") {
  339. inImports = true
  340. }
  341. if inImports && (strings.HasPrefix(s, "var") ||
  342. strings.HasPrefix(s, "func") ||
  343. strings.HasPrefix(s, "const") ||
  344. strings.HasPrefix(s, "type")) {
  345. done = true
  346. inImports = false
  347. }
  348. if inImports && len(breaks) > 0 {
  349. if m := impLine.FindStringSubmatch(s); m != nil {
  350. if m[1] == breaks[0] {
  351. out.WriteByte('\n')
  352. breaks = breaks[1:]
  353. }
  354. }
  355. }
  356. fmt.Fprint(&out, s)
  357. }
  358. return out.Bytes(), nil
  359. }