github.com/profzone/eden-framework@v1.0.10/internal/generator/format/format.go (about)

     1  package format
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"go/ast"
     7  	"go/format"
     8  	"go/parser"
     9  	"go/token"
    10  	"golang.org/x/tools/go/packages"
    11  	"os"
    12  	"runtime"
    13  	"strconv"
    14  	"strings"
    15  )
    16  
    17  type ImportGroups [][]string
    18  
    19  func getAstString(fileSet *token.FileSet, node ast.Node) string {
    20  	buf := &bytes.Buffer{}
    21  	if err := format.Node(buf, fileSet, node); err != nil {
    22  		panic(err)
    23  	}
    24  	return buf.String()
    25  }
    26  
    27  func Format(filename string, src []byte) []byte {
    28  	fileSet := token.NewFileSet()
    29  	file, err := parser.ParseFile(fileSet, filename, src, parser.ParseComments)
    30  	if err != nil {
    31  		panic(fmt.Errorf("errors %s in %s", err.Error(), filename))
    32  	}
    33  	buf := &bytes.Buffer{}
    34  	if err := format.Node(buf, fileSet, file); err != nil {
    35  		panic(fmt.Errorf("errors %s in %s", err.Error(), filename))
    36  	}
    37  	return buf.Bytes()
    38  }
    39  
    40  func Process(filename string, src []byte) ([]byte, error) {
    41  	cwd, _ := os.Getwd()
    42  	fileSet := token.NewFileSet()
    43  	file, err := parser.ParseFile(fileSet, filename, src, parser.ParseComments)
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  
    48  	ast.SortImports(fileSet, file)
    49  
    50  	formattedCode := getAstString(fileSet, file)
    51  
    52  	for _, decl := range file.Decls {
    53  		if genDecl, ok := decl.(*ast.GenDecl); ok {
    54  			if genDecl.Tok != token.IMPORT {
    55  				break
    56  			}
    57  
    58  			importsCode := getAstString(fileSet, genDecl)
    59  
    60  			importGroups := make(ImportGroups, 4)
    61  			for _, spec := range genDecl.Specs {
    62  				importSpec := spec.(*ast.ImportSpec)
    63  				importPath, _ := strconv.Unquote(importSpec.Path.Value)
    64  				pkgs, err := packages.Load(nil, importPath)
    65  				if err != nil {
    66  					panic(fmt.Errorf("errors %s in %s", err.Error(), filename))
    67  				}
    68  				errCount := packages.PrintErrors(pkgs)
    69  				if errCount > 0 {
    70  					panic(fmt.Errorf("you've got %d errors when packages.Load", errCount))
    71  				}
    72  				if strings.Contains(pkgs[0].GoFiles[0], runtime.GOROOT()) {
    73  					// libexec
    74  					importGroups[0] = append(importGroups[0], getAstString(fileSet, importSpec))
    75  				} else {
    76  					if strings.HasPrefix(pkgs[0].GoFiles[0], cwd) {
    77  						importGroups[3] = append(importGroups[3], getAstString(fileSet, importSpec))
    78  					} else {
    79  						if strings.HasPrefix(pkgs[0].PkgPath, "profzone") {
    80  							importGroups[2] = append(importGroups[2], getAstString(fileSet, importSpec))
    81  						} else {
    82  							importGroups[1] = append(importGroups[1], getAstString(fileSet, importSpec))
    83  						}
    84  					}
    85  				}
    86  			}
    87  
    88  			buf := &bytes.Buffer{}
    89  
    90  			buf.WriteString("import (\n")
    91  			for _, importGroup := range importGroups {
    92  				for _, code := range importGroup {
    93  					buf.WriteString(code + "\n")
    94  				}
    95  				buf.WriteString("\n")
    96  			}
    97  			buf.WriteString(")")
    98  			formattedCode = strings.Replace(formattedCode, importsCode, buf.String(), -1)
    99  		}
   100  	}
   101  
   102  	return Format(filename, []byte(formattedCode)), nil
   103  }