libgo/go/cmd/gofmt/gofmt.go

   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         // debugging
  35         cpuprofile = flag.String("cpuprofile", "", "write cpu profile to this file")
  36 )
  37
  38 const (
  39         tabWidth    = 8
  40         printerMode = printer.UseSpaces | printer.TabIndent
  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 )
  49
  50 func report(err error) {
  51         scanner.PrintError(os.Stderr, err)
  52         exitCode = 2
  53 }
  54
  55 func usage() {
  56         fmt.Fprintf(os.Stderr, "usage: gofmt [flags] [path ...]\n")
  57         flag.PrintDefaults()
  58         os.Exit(2)
  59 }
  60
  61 func initParserMode() {
  62         parserMode = parser.ParseComments
  63         if *allErrors {
  64                 parserMode |= parser.AllErrors
  65         }
  66 }
  67
  68 func isGoFile(f os.FileInfo) bool {
  69         // ignore non-Go files
  70         name := f.Name()
  71         return !f.IsDir() && !strings.HasPrefix(name, ".") && strings.HasSuffix(name, ".go")
  72 }
  73
  74 // If in == nil, the source is the contents of the file with the given filename.
  75 func processFile(filename string, in io.Reader, out io.Writer, stdin bool) error {
  76         if in == nil {
  77                 f, err := os.Open(filename)
  78                 if err != nil {
  79                         return err
  80                 }
  81                 defer f.Close()
  82                 in = f
  83         }
  84
  85         src, err := ioutil.ReadAll(in)
  86         if err != nil {
  87                 return err
  88         }
  89
  90         file, sourceAdj, indentAdj, err := parse(fileSet, filename, src, stdin)
  91         if err != nil {
  92                 return err
  93         }
  94
  95         if rewrite != nil {
  96                 if sourceAdj == nil {
  97                         file = rewrite(file)
  98                 } else {
  99                         fmt.Fprintf(os.Stderr, "warning: rewrite ignored for incomplete programs\n")
 100                 }
 101         }
 102
 103         ast.SortImports(fileSet, file)
 104
 105         if *simplifyAST {
 106                 simplify(file)
 107         }
 108
 109         res, err := format(fileSet, file, sourceAdj, indentAdj, src, printer.Config{Mode: printerMode, Tabwidth: tabWidth})
 110         if err != nil {
 111                 return err
 112         }
 113
 114         if !bytes.Equal(src, res) {
 115                 // formatting has changed
 116                 if *list {
 117                         fmt.Fprintln(out, filename)
 118                 }
 119                 if *write {
 120                         err = ioutil.WriteFile(filename, res, 0644)
 121                         if err != nil {
 122                                 return err
 123                         }
 124                 }
 125                 if *doDiff {
 126                         data, err := diff(src, res)
 127                         if err != nil {
 128                                 return fmt.Errorf("computing diff: %s", err)
 129                         }
 130                         fmt.Printf("diff %s gofmt/%s\n", filename, filename)
 131                         out.Write(data)
 132                 }
 133         }
 134
 135         if !*list && !*write && !*doDiff {
 136                 _, err = out.Write(res)
 137         }
 138
 139         return err
 140 }
 141
 142 func visitFile(path string, f os.FileInfo, err error) error {
 143         if err == nil && isGoFile(f) {
 144                 err = processFile(path, nil, os.Stdout, false)
 145         }
 146         if err != nil {
 147                 report(err)
 148         }
 149         return nil
 150 }
 151
 152 func walkDir(path string) {
 153         filepath.Walk(path, visitFile)
 154 }
 155
 156 func main() {
 157         // call gofmtMain in a separate function
 158         // so that it can use defer and have them
 159         // run before the exit.
 160         gofmtMain()
 161         os.Exit(exitCode)
 162 }
 163
 164 func gofmtMain() {
 165         flag.Usage = usage
 166         flag.Parse()
 167
 168         if *cpuprofile != "" {
 169                 f, err := os.Create(*cpuprofile)
 170                 if err != nil {
 171                         fmt.Fprintf(os.Stderr, "creating cpu profile: %s\n", err)
 172                         exitCode = 2
 173                         return
 174                 }
 175                 defer f.Close()
 176                 pprof.StartCPUProfile(f)
 177                 defer pprof.StopCPUProfile()
 178         }
 179
 180         initParserMode()
 181         initRewrite()
 182
 183         if flag.NArg() == 0 {
 184                 if *write {
 185                         fmt.Fprintln(os.Stderr, "error: cannot use -w with standard input")
 186                         exitCode = 2
 187                         return
 188                 }
 189                 if err := processFile("<standard input>", os.Stdin, os.Stdout, true); err != nil {
 190                         report(err)
 191                 }
 192                 return
 193         }
 194
 195         for i := 0; i < flag.NArg(); i++ {
 196                 path := flag.Arg(i)
 197                 switch dir, err := os.Stat(path); {
 198                 case err != nil:
 199                         report(err)
 200                 case dir.IsDir():
 201                         walkDir(path)
 202                 default:
 203                         if err := processFile(path, nil, os.Stdout, false); err != nil {
 204                                 report(err)
 205                         }
 206                 }
 207         }
 208 }
 209
 210 func diff(b1, b2 []byte) (data []byte, err error) {
 211         f1, err := ioutil.TempFile("", "gofmt")
 212         if err != nil {
 213                 return
 214         }
 215         defer os.Remove(f1.Name())
 216         defer f1.Close()
 217
 218         f2, err := ioutil.TempFile("", "gofmt")
 219         if err != nil {
 220                 return
 221         }
 222         defer os.Remove(f2.Name())
 223         defer f2.Close()
 224
 225         f1.Write(b1)
 226         f2.Write(b2)
 227
 228         data, err = exec.Command("diff", "-u", f1.Name(), f2.Name()).CombinedOutput()
 229         if len(data) > 0 {
 230                 // diff exits with a non-zero status when the files don't match.
 231                 // Ignore that failure as long as we get output.
 232                 err = nil
 233         }
 234         return
 235
 236 }
 237
 238 // ----------------------------------------------------------------------------
 239 // Support functions
 240 //
 241 // The functions parse, format, and isSpace below are identical to the
 242 // respective functions in src/go/format/format.go - keep them in sync!
 243 //
 244 // TODO(gri) Factor out this functionality, eventually.
 245
 246 // parse parses src, which was read from the named file,
 247 // as a Go source file, declaration, or statement list.
 248 func parse(fset *token.FileSet, filename string, src []byte, fragmentOk bool) (
 249         file *ast.File,
 250         sourceAdj func(src []byte, indent int) []byte,
 251         indentAdj int,
 252         err error,
 253 ) {
 254         // Try as whole source file.
 255         file, err = parser.ParseFile(fset, filename, src, parserMode)
 256         // If there's no error, return.  If the error is that the source file didn't begin with a
 257         // package line and source fragments are ok, fall through to
 258         // try as a source fragment.  Stop and return on any other error.
 259         if err == nil || !fragmentOk || !strings.Contains(err.Error(), "expected 'package'") {
 260                 return
 261         }
 262
 263         // If this is a declaration list, make it a source file
 264         // by inserting a package clause.
 265         // Insert using a ;, not a newline, so that the line numbers
 266         // in psrc match the ones in src.
 267         psrc := append([]byte("package p;"), src...)
 268         file, err = parser.ParseFile(fset, filename, psrc, parserMode)
 269         if err == nil {
 270                 sourceAdj = func(src []byte, indent int) []byte {
 271                         // Remove the package clause.
 272                         // Gofmt has turned the ; into a \n.
 273                         src = src[indent+len("package p\n"):]
 274                         return bytes.TrimSpace(src)
 275                 }
 276                 return
 277         }
 278         // If the error is that the source file didn't begin with a
 279         // declaration, fall through to try as a statement list.
 280         // Stop and return on any other error.
 281         if !strings.Contains(err.Error(), "expected declaration") {
 282                 return
 283         }
 284
 285         // If this is a statement list, make it a source file
 286         // by inserting a package clause and turning the list
 287         // into a function body.  This handles expressions too.
 288         // Insert using a ;, not a newline, so that the line numbers
 289         // in fsrc match the ones in src.
 290         fsrc := append(append([]byte("package p; func _() {"), src...), '\n', '}')
 291         file, err = parser.ParseFile(fset, filename, fsrc, parserMode)
 292         if err == nil {
 293                 sourceAdj = func(src []byte, indent int) []byte {
 294                         // Cap adjusted indent to zero.
 295                         if indent < 0 {
 296                                 indent = 0
 297                         }
 298                         // Remove the wrapping.
 299                         // Gofmt has turned the ; into a \n\n.
 300                         // There will be two non-blank lines with indent, hence 2*indent.
 301                         src = src[2*indent+len("package p\n\nfunc _() {"):]
 302                         src = src[:len(src)-(indent+len("\n}\n"))]
 303                         return bytes.TrimSpace(src)
 304                 }
 305                 // Gofmt has also indented the function body one level.
 306                 // Adjust that with indentAdj.
 307                 indentAdj = -1
 308         }
 309
 310         // Succeeded, or out of options.
 311         return
 312 }
 313
 314 // format formats the given package file originally obtained from src
 315 // and adjusts the result based on the original source via sourceAdj
 316 // and indentAdj.
 317 func format(
 318         fset *token.FileSet,
 319         file *ast.File,
 320         sourceAdj func(src []byte, indent int) []byte,
 321         indentAdj int,
 322         src []byte,
 323         cfg printer.Config,
 324 ) ([]byte, error) {
 325         if sourceAdj == nil {
 326                 // Complete source file.
 327                 var buf bytes.Buffer
 328                 err := cfg.Fprint(&buf, fset, file)
 329                 if err != nil {
 330                         return nil, err
 331                 }
 332                 return buf.Bytes(), nil
 333         }
 334
 335         // Partial source file.
 336         // Determine and prepend leading space.
 337         i, j := 0, 0
 338         for j < len(src) && isSpace(src[j]) {
 339                 if src[j] == '\n' {
 340                         i = j + 1 // byte offset of last line in leading space
 341                 }
 342                 j++
 343         }
 344         var res []byte
 345         res = append(res, src[:i]...)
 346
 347         // Determine and prepend indentation of first code line.
 348         // Spaces are ignored unless there are no tabs,
 349         // in which case spaces count as one tab.
 350         indent := 0
 351         hasSpace := false
 352         for _, b := range src[i:j] {
 353                 switch b {
 354                 case ' ':
 355                         hasSpace = true
 356                 case '\t':
 357                         indent++
 358                 }
 359         }
 360         if indent == 0 && hasSpace {
 361                 indent = 1
 362         }
 363         for i := 0; i < indent; i++ {
 364                 res = append(res, '\t')
 365         }
 366
 367         // Format the source.
 368         // Write it without any leading and trailing space.
 369         cfg.Indent = indent + indentAdj
 370         var buf bytes.Buffer
 371         err := cfg.Fprint(&buf, fset, file)
 372         if err != nil {
 373                 return nil, err
 374         }
 375         res = append(res, sourceAdj(buf.Bytes(), cfg.Indent)...)
 376
 377         // Determine and append trailing space.
 378         i = len(src)
 379         for i > 0 && isSpace(src[i-1]) {
 380                 i--
 381         }
 382         return append(res, src[i:]...), nil
 383 }
 384
 385 func isSpace(b byte) bool {
 386         return b == ' ' || b == '\t' || b == '\n' || b == '\r'
 387 }