goimports.go

Documentation: golang.org/x/tools/cmd/goimports

     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 main
     6  
     7  import (
     8  	"bufio"
     9  	"bytes"
    10  	"errors"
    11  	"flag"
    12  	"fmt"
    13  	"go/scanner"
    14  	"io"
    15  	"log"
    16  	"os"
    17  	"os/exec"
    18  	"path/filepath"
    19  	"runtime"
    20  	"runtime/pprof"
    21  	"strings"
    22  
    23  	"golang.org/x/tools/internal/gocommand"
    24  	"golang.org/x/tools/internal/imports"
    25  )
    26  
    27  var (
    28  	// main operation modes
    29  	list   = flag.Bool("l", false, "list files whose formatting differs from goimport's")
    30  	write  = flag.Bool("w", false, "write result to (source) file instead of stdout")
    31  	doDiff = flag.Bool("d", false, "display diffs instead of rewriting files")
    32  	srcdir = flag.String("srcdir", "", "choose imports as if source code is from `dir`. When operating on a single file, dir may instead be the complete file name.")
    33  
    34  	verbose bool // verbose logging
    35  
    36  	cpuProfile     = flag.String("cpuprofile", "", "CPU profile output")
    37  	memProfile     = flag.String("memprofile", "", "memory profile output")
    38  	memProfileRate = flag.Int("memrate", 0, "if > 0, sets runtime.MemProfileRate")
    39  
    40  	options = &imports.Options{
    41  		TabWidth:  8,
    42  		TabIndent: true,
    43  		Comments:  true,
    44  		Fragment:  true,
    45  		Env: &imports.ProcessEnv{
    46  			GocmdRunner: &gocommand.Runner{},
    47  		},
    48  	}
    49  	exitCode = 0
    50  )
    51  
    52  func init() {
    53  	flag.BoolVar(&options.AllErrors, "e", false, "report all errors (not just the first 10 on different lines)")
    54  	flag.StringVar(&options.LocalPrefix, "local", "", "put imports beginning with this string after 3rd-party packages; comma-separated list")
    55  	flag.BoolVar(&options.FormatOnly, "format-only", false, "if true, don't fix imports and only format. In this mode, goimports is effectively gofmt, with the addition that imports are grouped into sections.")
    56  }
    57  
    58  func report(err error) {
    59  	scanner.PrintError(os.Stderr, err)
    60  	exitCode = 2
    61  }
    62  
    63  func usage() {
    64  	fmt.Fprintf(os.Stderr, "usage: goimports [flags] [path ...]\n")
    65  	flag.PrintDefaults()
    66  	os.Exit(2)
    67  }
    68  
    69  func isGoFile(f os.FileInfo) bool {
    70  	// ignore non-Go files
    71  	name := f.Name()
    72  	return !f.IsDir() && !strings.HasPrefix(name, ".") && strings.HasSuffix(name, ".go")
    73  }
    74  
    75  // argumentType is which mode goimports was invoked as.
    76  type argumentType int
    77  
    78  const (
    79  	// fromStdin means the user is piping their source into goimports.
    80  	fromStdin argumentType = iota
    81  
    82  	// singleArg is the common case from editors, when goimports is run on
    83  	// a single file.
    84  	singleArg
    85  
    86  	// multipleArg is when the user ran "goimports file1.go file2.go"
    87  	// or ran goimports on a directory tree.
    88  	multipleArg
    89  )
    90  
    91  func processFile(filename string, in io.Reader, out io.Writer, argType argumentType) error {
    92  	opt := options
    93  	if argType == fromStdin {
    94  		nopt := *options
    95  		nopt.Fragment = true
    96  		opt = &nopt
    97  	}
    98  
    99  	if in == nil {
   100  		f, err := os.Open(filename)
   101  		if err != nil {
   102  			return err
   103  		}
   104  		defer f.Close()
   105  		in = f
   106  	}
   107  
   108  	src, err := io.ReadAll(in)
   109  	if err != nil {
   110  		return err
   111  	}
   112  
   113  	target := filename
   114  	if *srcdir != "" {
   115  		// Determine whether the provided -srcdirc is a directory or file
   116  		// and then use it to override the target.
   117  		//
   118  		// See https://github.com/dominikh/go-mode.el/issues/146
   119  		if isFile(*srcdir) {
   120  			if argType == multipleArg {
   121  				return errors.New("-srcdir value can't be a file when passing multiple arguments or when walking directories")
   122  			}
   123  			target = *srcdir
   124  		} else if argType == singleArg && strings.HasSuffix(*srcdir, ".go") && !isDir(*srcdir) {
   125  			// For a file which doesn't exist on disk yet, but might shortly.
   126  			// e.g. user in editor opens $DIR/newfile.go and newfile.go doesn't yet exist on disk.
   127  			// The goimports on-save hook writes the buffer to a temp file
   128  			// first and runs goimports before the actual save to newfile.go.
   129  			// The editor's buffer is named "newfile.go" so that is passed to goimports as:
   130  			//      goimports -srcdir=/gopath/src/pkg/newfile.go /tmp/gofmtXXXXXXXX.go
   131  			// and then the editor reloads the result from the tmp file and writes
   132  			// it to newfile.go.
   133  			target = *srcdir
   134  		} else {
   135  			// Pretend that file is from *srcdir in order to decide
   136  			// visible imports correctly.
   137  			target = filepath.Join(*srcdir, filepath.Base(filename))
   138  		}
   139  	}
   140  
   141  	res, err := imports.Process(target, src, opt)
   142  	if err != nil {
   143  		return err
   144  	}
   145  
   146  	if !bytes.Equal(src, res) {
   147  		// formatting has changed
   148  		if *list {
   149  			fmt.Fprintln(out, filename)
   150  		}
   151  		if *write {
   152  			if argType == fromStdin {
   153  				// filename is "<standard input>"
   154  				return errors.New("can't use -w on stdin")
   155  			}
   156  			// On Windows, we need to re-set the permissions from the file. See golang/go#38225.
   157  			var perms os.FileMode
   158  			if fi, err := os.Stat(filename); err == nil {
   159  				perms = fi.Mode() & os.ModePerm
   160  			}
   161  			err = os.WriteFile(filename, res, perms)
   162  			if err != nil {
   163  				return err
   164  			}
   165  		}
   166  		if *doDiff {
   167  			if argType == fromStdin {
   168  				filename = "stdin.go" // because <standard input>.orig looks silly
   169  			}
   170  			data, err := diff(src, res, filename)
   171  			if err != nil {
   172  				return fmt.Errorf("computing diff: %s", err)
   173  			}
   174  			fmt.Printf("diff -u %s %s\n", filepath.ToSlash(filename+".orig"), filepath.ToSlash(filename))
   175  			out.Write(data)
   176  		}
   177  	}
   178  
   179  	if !*list && !*write && !*doDiff {
   180  		_, err = out.Write(res)
   181  	}
   182  
   183  	return err
   184  }
   185  
   186  func visitFile(path string, f os.FileInfo, err error) error {
   187  	if err == nil && isGoFile(f) {
   188  		err = processFile(path, nil, os.Stdout, multipleArg)
   189  	}
   190  	if err != nil {
   191  		report(err)
   192  	}
   193  	return nil
   194  }
   195  
   196  func walkDir(path string) {
   197  	filepath.Walk(path, visitFile)
   198  }
   199  
   200  func main() {
   201  	runtime.GOMAXPROCS(runtime.NumCPU())
   202  
   203  	// call gofmtMain in a separate function
   204  	// so that it can use defer and have them
   205  	// run before the exit.
   206  	gofmtMain()
   207  	os.Exit(exitCode)
   208  }
   209  
   210  // parseFlags parses command line flags and returns the paths to process.
   211  // It's a var so that custom implementations can replace it in other files.
   212  var parseFlags = func() []string {
   213  	flag.BoolVar(&verbose, "v", false, "verbose logging")
   214  
   215  	flag.Parse()
   216  	return flag.Args()
   217  }
   218  
   219  func bufferedFileWriter(dest string) (w io.Writer, close func()) {
   220  	f, err := os.Create(dest)
   221  	if err != nil {
   222  		log.Fatal(err)
   223  	}
   224  	bw := bufio.NewWriter(f)
   225  	return bw, func() {
   226  		if err := bw.Flush(); err != nil {
   227  			log.Fatalf("error flushing %v: %v", dest, err)
   228  		}
   229  		if err := f.Close(); err != nil {
   230  			log.Fatal(err)
   231  		}
   232  	}
   233  }
   234  
   235  func gofmtMain() {
   236  	flag.Usage = usage
   237  	paths := parseFlags()
   238  
   239  	if *cpuProfile != "" {
   240  		bw, flush := bufferedFileWriter(*cpuProfile)
   241  		pprof.StartCPUProfile(bw)
   242  		defer flush()
   243  		defer pprof.StopCPUProfile()
   244  	}
   245  	// doTrace is a conditionally compiled wrapper around runtime/trace. It is
   246  	// used to allow goimports to compile under gccgo, which does not support
   247  	// runtime/trace. See https://golang.org/issue/15544.
   248  	defer doTrace()()
   249  	if *memProfileRate > 0 {
   250  		runtime.MemProfileRate = *memProfileRate
   251  		bw, flush := bufferedFileWriter(*memProfile)
   252  		defer func() {
   253  			runtime.GC() // materialize all statistics
   254  			if err := pprof.WriteHeapProfile(bw); err != nil {
   255  				log.Fatal(err)
   256  			}
   257  			flush()
   258  		}()
   259  	}
   260  
   261  	if verbose {
   262  		log.SetFlags(log.LstdFlags | log.Lmicroseconds)
   263  		options.Env.Logf = log.Printf
   264  	}
   265  	if options.TabWidth < 0 {
   266  		fmt.Fprintf(os.Stderr, "negative tabwidth %d\n", options.TabWidth)
   267  		exitCode = 2
   268  		return
   269  	}
   270  
   271  	if len(paths) == 0 {
   272  		if err := processFile("<standard input>", os.Stdin, os.Stdout, fromStdin); err != nil {
   273  			report(err)
   274  		}
   275  		return
   276  	}
   277  
   278  	argType := singleArg
   279  	if len(paths) > 1 {
   280  		argType = multipleArg
   281  	}
   282  
   283  	for _, path := range paths {
   284  		switch dir, err := os.Stat(path); {
   285  		case err != nil:
   286  			report(err)
   287  		case dir.IsDir():
   288  			walkDir(path)
   289  		default:
   290  			if err := processFile(path, nil, os.Stdout, argType); err != nil {
   291  				report(err)
   292  			}
   293  		}
   294  	}
   295  }
   296  
   297  func writeTempFile(dir, prefix string, data []byte) (string, error) {
   298  	file, err := os.CreateTemp(dir, prefix)
   299  	if err != nil {
   300  		return "", err
   301  	}
   302  	_, err = file.Write(data)
   303  	if err1 := file.Close(); err == nil {
   304  		err = err1
   305  	}
   306  	if err != nil {
   307  		os.Remove(file.Name())
   308  		return "", err
   309  	}
   310  	return file.Name(), nil
   311  }
   312  
   313  func diff(b1, b2 []byte, filename string) (data []byte, err error) {
   314  	f1, err := writeTempFile("", "gofmt", b1)
   315  	if err != nil {
   316  		return
   317  	}
   318  	defer os.Remove(f1)
   319  
   320  	f2, err := writeTempFile("", "gofmt", b2)
   321  	if err != nil {
   322  		return
   323  	}
   324  	defer os.Remove(f2)
   325  
   326  	cmd := "diff"
   327  	if runtime.GOOS == "plan9" {
   328  		cmd = "/bin/ape/diff"
   329  	}
   330  
   331  	data, err = exec.Command(cmd, "-u", f1, f2).CombinedOutput()
   332  	if len(data) > 0 {
   333  		// diff exits with a non-zero status when the files don't match.
   334  		// Ignore that failure as long as we get output.
   335  		return replaceTempFilename(data, filename)
   336  	}
   337  	return
   338  }
   339  
   340  // replaceTempFilename replaces temporary filenames in diff with actual one.
   341  //
   342  // --- /tmp/gofmt316145376	2017-02-03 19:13:00.280468375 -0500
   343  // +++ /tmp/gofmt617882815	2017-02-03 19:13:00.280468375 -0500
   344  // ...
   345  // ->
   346  // --- path/to/file.go.orig	2017-02-03 19:13:00.280468375 -0500
   347  // +++ path/to/file.go	2017-02-03 19:13:00.280468375 -0500
   348  // ...
   349  func replaceTempFilename(diff []byte, filename string) ([]byte, error) {
   350  	bs := bytes.SplitN(diff, []byte{'\n'}, 3)
   351  	if len(bs) < 3 {
   352  		return nil, fmt.Errorf("got unexpected diff for %s", filename)
   353  	}
   354  	// Preserve timestamps.
   355  	var t0, t1 []byte
   356  	if i := bytes.LastIndexByte(bs[0], '\t'); i != -1 {
   357  		t0 = bs[0][i:]
   358  	}
   359  	if i := bytes.LastIndexByte(bs[1], '\t'); i != -1 {
   360  		t1 = bs[1][i:]
   361  	}
   362  	// Always print filepath with slash separator.
   363  	f := filepath.ToSlash(filename)
   364  	bs[0] = []byte(fmt.Sprintf("--- %s%s", f+".orig", t0))
   365  	bs[1] = []byte(fmt.Sprintf("+++ %s%s", f, t1))
   366  	return bytes.Join(bs, []byte{'\n'}), nil
   367  }
   368  
   369  // isFile reports whether name is a file.
   370  func isFile(name string) bool {
   371  	fi, err := os.Stat(name)
   372  	return err == nil && fi.Mode().IsRegular()
   373  }
   374  
   375  // isDir reports whether name is a directory.
   376  func isDir(name string) bool {
   377  	fi, err := os.Stat(name)
   378  	return err == nil && fi.IsDir()
   379  }
   380
View as plain text