github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/gen/codegen/formatx/pkg_sort.go (about)

     1  package formatx
     2  
     3  import (
     4  	"bytes"
     5  	"go/ast"
     6  	"go/format"
     7  	"go/token"
     8  	"io/ioutil"
     9  	"path/filepath"
    10  	"runtime"
    11  	"strconv"
    12  	"strings"
    13  )
    14  
    15  func SortImports(fset *token.FileSet, f *ast.File, file, group string) error {
    16  	ast.SortImports(fset, f)
    17  	dir := filepath.Dir(file)
    18  
    19  	for _, decl := range f.Decls {
    20  		d, ok := decl.(*ast.GenDecl)
    21  		if !ok || d.Tok != token.IMPORT || len(d.Specs) == 0 {
    22  			break
    23  		}
    24  		g := &GroupSet{}
    25  		for i := range d.Specs {
    26  			s := d.Specs[i].(*ast.ImportSpec)
    27  			path, _ := strconv.Unquote(s.Path.Value)
    28  
    29  			if Stds[path] {
    30  				g.AppendStd(path, s)
    31  				continue
    32  			}
    33  			if group != "" && strings.HasPrefix(path, group) {
    34  				g.AppendGroup(path, s)
    35  				continue
    36  			}
    37  			if strings.Contains(dir, path) {
    38  				g.AppendLocal(path, s)
    39  				continue
    40  			}
    41  			g.AppendVendor(path, s)
    42  		}
    43  		_fset, _f, err := Parse(file, bytes.Replace(
    44  			FmtNode(fset, f),
    45  			FmtNode(fset, d),
    46  			g.Bytes(),
    47  			1,
    48  		))
    49  		if err != nil {
    50  			return err
    51  		}
    52  		*fset, *f = *_fset, *_f
    53  	}
    54  
    55  	return nil
    56  }
    57  
    58  func FmtNode(fset *token.FileSet, node ast.Node) []byte {
    59  	buf := bytes.NewBuffer(nil)
    60  	if err := format.Node(buf, fset, node); err != nil {
    61  		panic(err)
    62  	}
    63  	return buf.Bytes()
    64  }
    65  
    66  type dep struct {
    67  	Pkg  string
    68  	Spec *ast.ImportSpec
    69  }
    70  
    71  type GroupSet [4][]*dep
    72  
    73  func (g GroupSet) Bytes() []byte {
    74  	buf := bytes.NewBuffer(nil)
    75  
    76  	buf.WriteString("import ")
    77  	if g.Len() > 1 {
    78  		buf.WriteRune('(')
    79  	}
    80  	for _, deps := range g {
    81  		for _, d := range deps {
    82  			buf.WriteRune('\n')
    83  			if d.Spec.Doc != nil {
    84  				for _, c := range d.Spec.Doc.List {
    85  					buf.WriteString(c.Text)
    86  					buf.WriteRune('\n')
    87  				}
    88  			}
    89  			if d.Spec.Name != nil {
    90  				buf.WriteString(d.Spec.Name.String())
    91  				buf.WriteRune(' ')
    92  			}
    93  			buf.WriteString(d.Spec.Path.Value)
    94  			if d.Spec.Comment != nil {
    95  				for _, c := range d.Spec.Comment.List {
    96  					buf.WriteString(c.Text)
    97  				}
    98  			}
    99  		}
   100  		buf.WriteRune('\n')
   101  	}
   102  	if g.Len() > 1 {
   103  		buf.WriteRune(')')
   104  	}
   105  
   106  	return buf.Bytes()
   107  }
   108  
   109  func (g *GroupSet) Len() int {
   110  	return len(g[0]) + len(g[1]) + len(g[2]) + len(g[3])
   111  }
   112  
   113  func (g *GroupSet) append(idx int, pkg string, spec *ast.ImportSpec) {
   114  	(*g)[idx] = append((*g)[idx], &dep{Pkg: pkg, Spec: spec})
   115  }
   116  
   117  func (g *GroupSet) AppendStd(pkg string, spec *ast.ImportSpec)    { g.append(0, pkg, spec) }
   118  func (g *GroupSet) AppendVendor(pkg string, spec *ast.ImportSpec) { g.append(1, pkg, spec) }
   119  func (g *GroupSet) AppendLocal(pkg string, spec *ast.ImportSpec)  { g.append(2, pkg, spec) }
   120  func (g *GroupSet) AppendGroup(pkg string, spec *ast.ImportSpec)  { g.append(3, pkg, spec) }
   121  
   122  type StdLibSet map[string]bool
   123  
   124  func (s StdLibSet) WalkInit(root, prefix string) {
   125  	ds, _ := ioutil.ReadDir(root)
   126  	for _, d := range ds {
   127  		if !d.IsDir() {
   128  			continue
   129  		}
   130  		name := d.Name()
   131  		if name == "vendor" || name == "internal" || name == "testdata" {
   132  			continue
   133  		}
   134  		pkg := name
   135  		if prefix != "" {
   136  			pkg = filepath.Join(prefix, pkg)
   137  		}
   138  		s.WalkInit(filepath.Join(root, name), pkg)
   139  		s[pkg] = true
   140  	}
   141  }
   142  
   143  var Stds StdLibSet
   144  
   145  func init() {
   146  	Stds = make(StdLibSet)
   147  	Stds.WalkInit(filepath.Join(runtime.GOROOT(), "src"), "")
   148  	// root := filepath.Join(runtime.GOROOT(), "src")
   149  	// err := filepath.Walk(
   150  	// 	root,
   151  	// 	func(path string, info fs.FileInfo, err error) error {
   152  	// 		if path == root {
   153  	// 			return nil
   154  	// 		}
   155  	// 		if err != nil {
   156  	// 			return err
   157  	// 		}
   158  	// 		if !info.IsDir() {
   159  	// 			return nil
   160  	// 		}
   161  	// 		Stds[path] = true
   162  	// 		return nil
   163  	// 	},
   164  	// )
   165  	// if err != nil {
   166  	// 	panic(err)
   167  	// }
   168  }