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