github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/kit/enumgen/generator.go (about)

     1  package enumgen
     2  
     3  import (
     4  	"go/types"
     5  	"log"
     6  	"path"
     7  	"path/filepath"
     8  	"sort"
     9  
    10  	"golang.org/x/tools/go/packages"
    11  
    12  	"github.com/machinefi/w3bstream/pkg/depends/gen/codegen"
    13  	"github.com/machinefi/w3bstream/pkg/depends/x/misc/must"
    14  	"github.com/machinefi/w3bstream/pkg/depends/x/pkgx"
    15  	"github.com/machinefi/w3bstream/pkg/depends/x/stringsx"
    16  )
    17  
    18  type Generator struct {
    19  	pkg *pkgx.Pkg
    20  	*Scanner
    21  	enums map[*types.TypeName]*Enum
    22  }
    23  
    24  func New(pkg *pkgx.Pkg) *Generator {
    25  	return &Generator{
    26  		pkg:     pkg,
    27  		Scanner: NewScanner(pkg),
    28  		enums:   map[*types.TypeName]*Enum{},
    29  	}
    30  }
    31  
    32  func (g *Generator) Scan(names ...string) {
    33  	for _, name := range names {
    34  		tn := g.pkg.TypeName(name)
    35  		if tn == nil {
    36  			continue
    37  		}
    38  		opts, ok := g.Scanner.Options(tn)
    39  		if ok && opts[0].Int != nil && opts[0].Str != nil {
    40  			sort.Slice(opts, func(i, j int) bool {
    41  				return *opts[i].Int < *opts[j].Int
    42  			})
    43  			g.enums[tn] = NewEnum(tn.Pkg().Path()+"."+tn.Name(), opts)
    44  		}
    45  	}
    46  }
    47  
    48  func (g Generator) Output(cwd string) {
    49  	for tn, enum := range g.enums {
    50  		dir, _ := filepath.Rel(
    51  			cwd,
    52  			must.String(pkgx.PkgPathByPath(tn.Pkg().Path(), packages.NeedName, packages.NeedFiles)),
    53  		)
    54  		filename := codegen.GenerateFileSuffix(
    55  			path.Join(dir, stringsx.LowerSnakeCase(enum.Name)+".go"))
    56  		f := codegen.NewFile(tn.Pkg().Name(), filename)
    57  		enum.WriteToFile(f)
    58  
    59  		if _, err := f.Write(); err != nil {
    60  			log.Printf("%s generate failed: %v", filename, err)
    61  		}
    62  	}
    63  }
    64  
    65  func GetEnumByName(g *Generator, name string) *Enum {
    66  	for tn, enum := range g.enums {
    67  		if tn.Name() == name {
    68  			return enum
    69  		}
    70  	}
    71  	return nil
    72  }