github.com/xushiwei/go@v0.0.0-20130601165731-2b9d83f45bc9/src/cmd/gofmt/gofmt.go (about)

     1  // Copyright 2009 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  	"bytes"
     9  	"flag"
    10  	"fmt"
    11  	"go/ast"
    12  	"go/parser"
    13  	"go/printer"
    14  	"go/scanner"
    15  	"go/token"
    16  	"io"
    17  	"io/ioutil"
    18  	"os"
    19  	"os/exec"
    20  	"path/filepath"
    21  	"runtime/pprof"
    22  	"strings"
    23  )
    24  
    25  var (
    26  	// main operation modes
    27  	list        = flag.Bool("l", false, "list files whose formatting differs from gofmt's")
    28  	write       = flag.Bool("w", false, "write result to (source) file instead of stdout")
    29  	rewriteRule = flag.String("r", "", "rewrite rule (e.g., 'a[b:len(a)] -> a[b:]')")
    30  	simplifyAST = flag.Bool("s", false, "simplify code")
    31  	doDiff      = flag.Bool("d", false, "display diffs instead of rewriting files")
    32  	allErrors   = flag.Bool("e", false, "report all errors (not just the first 10 on different lines)")
    33  
    34  	// layout control
    35  	comments  = flag.Bool("comments", true, "print comments")
    36  	tabWidth  = flag.Int("tabwidth", 8, "tab width")
    37  	tabIndent = flag.Bool("tabs", true, "indent with tabs")
    38  
    39  	// debugging
    40  	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to this file")
    41  )
    42  
    43  var (
    44  	fileSet     = token.NewFileSet() // per process FileSet
    45  	exitCode    = 0
    46  	rewrite     func(*ast.File) *ast.File
    47  	parserMode  parser.Mode
    48  	printerMode printer.Mode
    49  )
    50  
    51  func report(err error) {
    52  	scanner.PrintError(os.Stderr, err)
    53  	exitCode = 2
    54  }
    55  
    56  func usage() {
    57  	fmt.Fprintf(os.Stderr, "usage: gofmt [flags] [path ...]\n")
    58  	flag.PrintDefaults()
    59  	os.Exit(2)
    60  }
    61  
    62  func initParserMode() {
    63  	parserMode = parser.Mode(0)
    64  	if *comments {
    65  		parserMode |= parser.ParseComments
    66  	}
    67  	if *allErrors {
    68  		parserMode |= parser.AllErrors
    69  	}
    70  }
    71  
    72  func initPrinterMode() {
    73  	printerMode = printer.UseSpaces
    74  	if *tabIndent {
    75  		printerMode |= printer.TabIndent
    76  	}
    77  }
    78  
    79  func isGoFile(f os.FileInfo) bool {
    80  	// ignore non-Go files
    81  	name := f.Name()
    82  	return !f.IsDir() && !strings.HasPrefix(name, ".") && strings.HasSuffix(name, ".go")
    83  }
    84  
    85  // If in == nil, the source is the contents of the file with the given filename.
    86  func processFile(filename string, in io.Reader, out io.Writer, stdin bool) error {
    87  	if in == nil {
    88  		f, err := os.Open(filename)
    89  		if err != nil {
    90  			return err
    91  		}
    92  		defer f.Close()
    93  		in = f
    94  	}
    95  
    96  	src, err := ioutil.ReadAll(in)
    97  	if err != nil {
    98  		return err
    99  	}
   100  
   101  	file, adjust, err := parse(fileSet, filename, src, stdin)
   102  	if err != nil {
   103  		return err
   104  	}
   105  
   106  	if rewrite != nil {
   107  		if adjust == nil {
   108  			file = rewrite(file)
   109  		} else {
   110  			fmt.Fprintf(os.Stderr, "warning: rewrite ignored for incomplete programs\n")
   111  		}
   112  	}
   113  
   114  	ast.SortImports(fileSet, file)
   115  
   116  	if *simplifyAST {
   117  		simplify(file)
   118  	}
   119  
   120  	var buf bytes.Buffer
   121  	err = (&printer.Config{Mode: printerMode, Tabwidth: *tabWidth}).Fprint(&buf, fileSet, file)
   122  	if err != nil {
   123  		return err
   124  	}
   125  	res := buf.Bytes()
   126  	if adjust != nil {
   127  		res = adjust(src, res)
   128  	}
   129  
   130  	if !bytes.Equal(src, res) {
   131  		// formatting has changed
   132  		if *list {
   133  			fmt.Fprintln(out, filename)
   134  		}
   135  		if *write {
   136  			err = ioutil.WriteFile(filename, res, 0)
   137  			if err != nil {
   138  				return err
   139  			}
   140  		}
   141  		if *doDiff {
   142  			data, err := diff(src, res)
   143  			if err != nil {
   144  				return fmt.Errorf("computing diff: %s", err)
   145  			}
   146  			fmt.Printf("diff %s gofmt/%s\n", filename, filename)
   147  			out.Write(data)
   148  		}
   149  	}
   150  
   151  	if !*list && !*write && !*doDiff {
   152  		_, err = out.Write(res)
   153  	}
   154  
   155  	return err
   156  }
   157  
   158  func visitFile(path string, f os.FileInfo, err error) error {
   159  	if err == nil && isGoFile(f) {
   160  		err = processFile(path, nil, os.Stdout, false)
   161  	}
   162  	if err != nil {
   163  		report(err)
   164  	}
   165  	return nil
   166  }
   167  
   168  func walkDir(path string) {
   169  	filepath.Walk(path, visitFile)
   170  }
   171  
   172  func main() {
   173  	// call gofmtMain in a separate function
   174  	// so that it can use defer and have them
   175  	// run before the exit.
   176  	gofmtMain()
   177  	os.Exit(exitCode)
   178  }
   179  
   180  func gofmtMain() {
   181  	flag.Usage = usage
   182  	flag.Parse()
   183  	if *tabWidth < 0 {
   184  		fmt.Fprintf(os.Stderr, "negative tabwidth %d\n", *tabWidth)
   185  		exitCode = 2
   186  		return
   187  	}
   188  
   189  	if *cpuprofile != "" {
   190  		f, err := os.Create(*cpuprofile)
   191  		if err != nil {
   192  			fmt.Fprintf(os.Stderr, "creating cpu profile: %s\n", err)
   193  			exitCode = 2
   194  			return
   195  		}
   196  		defer f.Close()
   197  		pprof.StartCPUProfile(f)
   198  		defer pprof.StopCPUProfile()
   199  	}
   200  
   201  	initParserMode()
   202  	initPrinterMode()
   203  	initRewrite()
   204  
   205  	if flag.NArg() == 0 {
   206  		if err := processFile("<standard input>", os.Stdin, os.Stdout, true); err != nil {
   207  			report(err)
   208  		}
   209  		return
   210  	}
   211  
   212  	for i := 0; i < flag.NArg(); i++ {
   213  		path := flag.Arg(i)
   214  		switch dir, err := os.Stat(path); {
   215  		case err != nil:
   216  			report(err)
   217  		case dir.IsDir():
   218  			walkDir(path)
   219  		default:
   220  			if err := processFile(path, nil, os.Stdout, false); err != nil {
   221  				report(err)
   222  			}
   223  		}
   224  	}
   225  }
   226  
   227  func diff(b1, b2 []byte) (data []byte, err error) {
   228  	f1, err := ioutil.TempFile("", "gofmt")
   229  	if err != nil {
   230  		return
   231  	}
   232  	defer os.Remove(f1.Name())
   233  	defer f1.Close()
   234  
   235  	f2, err := ioutil.TempFile("", "gofmt")
   236  	if err != nil {
   237  		return
   238  	}
   239  	defer os.Remove(f2.Name())
   240  	defer f2.Close()
   241  
   242  	f1.Write(b1)
   243  	f2.Write(b2)
   244  
   245  	data, err = exec.Command("diff", "-u", f1.Name(), f2.Name()).CombinedOutput()
   246  	if len(data) > 0 {
   247  		// diff exits with a non-zero status when the files don't match.
   248  		// Ignore that failure as long as we get output.
   249  		err = nil
   250  	}
   251  	return
   252  
   253  }
   254  
   255  // parse parses src, which was read from filename,
   256  // as a Go source file or statement list.
   257  func parse(fset *token.FileSet, filename string, src []byte, stdin bool) (*ast.File, func(orig, src []byte) []byte, error) {
   258  	// Try as whole source file.
   259  	file, err := parser.ParseFile(fset, filename, src, parserMode)
   260  	if err == nil {
   261  		return file, nil, nil
   262  	}
   263  	// If the error is that the source file didn't begin with a
   264  	// package line and this is standard input, fall through to
   265  	// try as a source fragment.  Stop and return on any other error.
   266  	if !stdin || !strings.Contains(err.Error(), "expected 'package'") {
   267  		return nil, nil, err
   268  	}
   269  
   270  	// If this is a declaration list, make it a source file
   271  	// by inserting a package clause.
   272  	// Insert using a ;, not a newline, so that the line numbers
   273  	// in psrc match the ones in src.
   274  	psrc := append([]byte("package p;"), src...)
   275  	file, err = parser.ParseFile(fset, filename, psrc, parserMode)
   276  	if err == nil {
   277  		adjust := func(orig, src []byte) []byte {
   278  			// Remove the package clause.
   279  			// Gofmt has turned the ; into a \n.
   280  			src = src[len("package p\n"):]
   281  			return matchSpace(orig, src)
   282  		}
   283  		return file, adjust, nil
   284  	}
   285  	// If the error is that the source file didn't begin with a
   286  	// declaration, fall through to try as a statement list.
   287  	// Stop and return on any other error.
   288  	if !strings.Contains(err.Error(), "expected declaration") {
   289  		return nil, nil, err
   290  	}
   291  
   292  	// If this is a statement list, make it a source file
   293  	// by inserting a package clause and turning the list
   294  	// into a function body.  This handles expressions too.
   295  	// Insert using a ;, not a newline, so that the line numbers
   296  	// in fsrc match the ones in src.
   297  	fsrc := append(append([]byte("package p; func _() {"), src...), '}')
   298  	file, err = parser.ParseFile(fset, filename, fsrc, parserMode)
   299  	if err == nil {
   300  		adjust := func(orig, src []byte) []byte {
   301  			// Remove the wrapping.
   302  			// Gofmt has turned the ; into a \n\n.
   303  			src = src[len("package p\n\nfunc _() {"):]
   304  			src = src[:len(src)-len("}\n")]
   305  			// Gofmt has also indented the function body one level.
   306  			// Remove that indent.
   307  			src = bytes.Replace(src, []byte("\n\t"), []byte("\n"), -1)
   308  			return matchSpace(orig, src)
   309  		}
   310  		return file, adjust, nil
   311  	}
   312  
   313  	// Failed, and out of options.
   314  	return nil, nil, err
   315  }
   316  
   317  func cutSpace(b []byte) (before, middle, after []byte) {
   318  	i := 0
   319  	for i < len(b) && (b[i] == ' ' || b[i] == '\t' || b[i] == '\n') {
   320  		i++
   321  	}
   322  	j := len(b)
   323  	for j > 0 && (b[j-1] == ' ' || b[j-1] == '\t' || b[j-1] == '\n') {
   324  		j--
   325  	}
   326  	if i <= j {
   327  		return b[:i], b[i:j], b[j:]
   328  	}
   329  	return nil, nil, b[j:]
   330  }
   331  
   332  // matchSpace reformats src to use the same space context as orig.
   333  // 1) If orig begins with blank lines, matchSpace inserts them at the beginning of src.
   334  // 2) matchSpace copies the indentation of the first non-blank line in orig
   335  //    to every non-blank line in src.
   336  // 3) matchSpace copies the trailing space from orig and uses it in place
   337  //   of src's trailing space.
   338  func matchSpace(orig []byte, src []byte) []byte {
   339  	before, _, after := cutSpace(orig)
   340  	i := bytes.LastIndex(before, []byte{'\n'})
   341  	before, indent := before[:i+1], before[i+1:]
   342  
   343  	_, src, _ = cutSpace(src)
   344  
   345  	var b bytes.Buffer
   346  	b.Write(before)
   347  	for len(src) > 0 {
   348  		line := src
   349  		if i := bytes.IndexByte(line, '\n'); i >= 0 {
   350  			line, src = line[:i+1], line[i+1:]
   351  		} else {
   352  			src = nil
   353  		}
   354  		if len(line) > 0 && line[0] != '\n' { // not blank
   355  			b.Write(indent)
   356  		}
   357  		b.Write(line)
   358  	}
   359  	b.Write(after)
   360  	return b.Bytes()
   361  }