github.com/sandwich-go/boost@v1.3.29/misc/goformat/goformat.go (about)

     1  package goformat
     2  
     3  import (
     4  	"bytes"
     5  	"go/ast"
     6  	"go/format"
     7  	"go/parser"
     8  	"go/printer"
     9  	"go/token"
    10  	"strings"
    11  )
    12  
    13  // ProcessFile
    14  // filename 源文件路径
    15  func ProcessFile(filename string, opts ...Option) ([]byte, error) {
    16  	return Process(filename, nil, opts...)
    17  }
    18  
    19  // ProcessCode
    20  // src 源码
    21  func ProcessCode(src []byte, opts ...Option) ([]byte, error) {
    22  	return Process("", src, opts...)
    23  }
    24  
    25  // Process 格式化文件或内容
    26  // filename 源文件路径
    27  // src 源码
    28  func Process(filename string, src []byte, opts ...Option) ([]byte, error) {
    29  	conf := NewOptions(opts...)
    30  
    31  	fileSet := token.NewFileSet()
    32  	file, adjust, err := parse(fileSet, filename, src, conf)
    33  	if err != nil {
    34  		return nil, err
    35  	}
    36  
    37  	if err = fillReturnValues(file); err != nil {
    38  		return nil, err
    39  	}
    40  
    41  	if conf.RemoveBareReturns {
    42  		if err = removeBareReturns(file); err != nil {
    43  			return nil, err
    44  		}
    45  	}
    46  
    47  	var buf bytes.Buffer
    48  	err = printer.Fprint(&buf, fileSet, file)
    49  	if err != nil {
    50  		return nil, err
    51  	}
    52  	out := buf.Bytes()
    53  	if adjust != nil {
    54  		out = adjust(src, out)
    55  	}
    56  	return format.Source(out)
    57  }
    58  
    59  func parseMainFragment(fset *token.FileSet, filename string, src []byte, mode parser.Mode) (*ast.File, func(orig, src []byte) []byte, error) {
    60  	psrc := append([]byte("package main;"), src...)
    61  	file, err := parser.ParseFile(fset, filename, psrc, mode)
    62  	if err == nil {
    63  		if containsMainFunc(file) {
    64  			return file, nil, nil
    65  		}
    66  		adjust := func(orig, src []byte) []byte {
    67  			src = src[len("package main\n"):]
    68  			return matchSpace(orig, src)
    69  		}
    70  		return file, adjust, nil
    71  	}
    72  	return nil, nil, err
    73  }
    74  
    75  func parseDeclarationFragment(fset *token.FileSet, filename string, src []byte, mode parser.Mode) (*ast.File, func(orig, src []byte) []byte, error) {
    76  	fsrc := append(append([]byte("package p; func _() {"), src...), '}')
    77  	file, err := parser.ParseFile(fset, filename, fsrc, mode)
    78  	if err == nil {
    79  		adjust := func(orig, src []byte) []byte {
    80  			src = src[len("package p\n\nfunc _() {"):]
    81  			src = src[:len(src)-len("}\n")]
    82  			src = bytes.Replace(src, []byte("\n\t"), []byte("\n"), -1)
    83  			return matchSpace(orig, src)
    84  		}
    85  		return file, adjust, nil
    86  	}
    87  	return nil, nil, err
    88  }
    89  
    90  func parse(fset *token.FileSet, filename string, src []byte, conf *Options) (*ast.File, func(orig, src []byte) []byte, error) {
    91  	mode := parser.ParseComments
    92  	if conf.AllErrors {
    93  		mode |= parser.AllErrors
    94  	}
    95  	var err error
    96  	var file *ast.File
    97  	if src == nil {
    98  		file, err = parser.ParseFile(fset, filename, nil, mode)
    99  	} else {
   100  		file, err = parser.ParseFile(fset, filename, src, mode)
   101  	}
   102  	if err == nil {
   103  		return file, nil, nil
   104  	}
   105  	if !conf.Fragment || !strings.Contains(err.Error(), "expected 'package'") {
   106  		return nil, nil, err
   107  	}
   108  
   109  	var adjust func(orig, src []byte) []byte
   110  	file, adjust, err = parseMainFragment(fset, filename, src, mode)
   111  	if err == nil {
   112  		return file, adjust, nil
   113  	}
   114  
   115  	if !strings.Contains(err.Error(), "expected declaration") {
   116  		return nil, nil, err
   117  	}
   118  	file, adjust, err = parseDeclarationFragment(fset, filename, src, mode)
   119  	if err == nil {
   120  		return file, adjust, nil
   121  	}
   122  	return nil, nil, err
   123  }