github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/courier/enumeration/gen/enum_generator.go (about)

     1  package gen
     2  
     3  import (
     4  	"go/build"
     5  	"go/parser"
     6  	"go/types"
     7  	"path"
     8  	"path/filepath"
     9  
    10  	"golang.org/x/tools/go/loader"
    11  
    12  	"github.com/johnnyeven/libtools/codegen"
    13  	"github.com/johnnyeven/libtools/codegen/loaderx"
    14  	"github.com/johnnyeven/libtools/courier/swagger/gen"
    15  	"github.com/johnnyeven/libtools/godash"
    16  )
    17  
    18  type EnumGenerator struct {
    19  	Filters       []string
    20  	pkgImportPath string
    21  	program       *loader.Program
    22  	enumScanner   *gen.EnumScanner
    23  }
    24  
    25  var _ interface {
    26  	codegen.Generator
    27  } = (*EnumGenerator)(nil)
    28  
    29  func (g *EnumGenerator) Load(cwd string) {
    30  	ldr := loader.Config{
    31  		AllowErrors: true,
    32  		ParserMode:  parser.ParseComments,
    33  	}
    34  
    35  	pkgImportPath := codegen.GetPackageImportPath(cwd)
    36  	ldr.Import(pkgImportPath)
    37  
    38  	p, err := ldr.Load()
    39  	if err != nil {
    40  		panic(err)
    41  	}
    42  
    43  	g.program = p
    44  	g.pkgImportPath = pkgImportPath
    45  	g.enumScanner = gen.NewEnumScanner(p)
    46  }
    47  
    48  func (g *EnumGenerator) Pick() {
    49  	for pkg, pkgInfo := range g.program.AllPackages {
    50  		if pkg.Path() != g.pkgImportPath {
    51  			continue
    52  		}
    53  		for ident, obj := range pkgInfo.Defs {
    54  			if typeName, ok := obj.(*types.TypeName); ok {
    55  				doc := loaderx.CommentsOf(g.program.Fset, ident, pkgInfo.Files...)
    56  				doc, hasEnum := gen.ParseEnum(doc)
    57  
    58  				if hasEnum {
    59  					if len(g.Filters) > 0 {
    60  						if godash.StringIncludes(g.Filters, typeName.Name()) {
    61  							g.enumScanner.Enum(typeName)
    62  						}
    63  					} else {
    64  						g.enumScanner.Enum(typeName)
    65  					}
    66  				}
    67  			}
    68  		}
    69  	}
    70  }
    71  
    72  func (g *EnumGenerator) Output(cwd string) codegen.Outputs {
    73  	outputs := codegen.Outputs{}
    74  	for typeName, e := range g.enumScanner.Enums {
    75  		p, _ := build.Import(typeName.Pkg().Path(), "", build.FindOnly)
    76  		dir, _ := filepath.Rel(cwd, p.Dir)
    77  
    78  		enum := NewEnum(typeName.Pkg().Path(), typeName.Pkg().Name(), typeName.Name(), e, g.enumScanner.HasOffset(typeName))
    79  
    80  		outputs.Add(codegen.GeneratedSuffix(path.Join(dir, codegen.ToLowerSnakeCase(typeName.Name())+".go")), enum.String())
    81  	}
    82  	return outputs
    83  }