github.com/unionj-cloud/go-doudou/v2@v2.3.5/toolkit/astutils/enumcollector.go (about)

     1  package astutils
     2  
     3  import (
     4  	"go/ast"
     5  	"go/parser"
     6  	"go/token"
     7  	"strings"
     8  )
     9  
    10  type EnumCollector struct {
    11  	Methods    map[string][]MethodMeta
    12  	Package    PackageMeta
    13  	exprString func(ast.Expr) string
    14  	Consts     map[string][]string
    15  }
    16  
    17  // Visit traverse each node from source code
    18  func (sc *EnumCollector) Visit(n ast.Node) ast.Visitor {
    19  	return sc.Collect(n)
    20  }
    21  
    22  // Collect collects all structs from source code
    23  func (sc *EnumCollector) Collect(n ast.Node) ast.Visitor {
    24  	switch spec := n.(type) {
    25  	case *ast.Package:
    26  		return sc
    27  	case *ast.File: // actually it is package name
    28  		sc.Package = PackageMeta{
    29  			Name: spec.Name.Name,
    30  		}
    31  		return sc
    32  	case *ast.FuncDecl:
    33  		if spec.Recv != nil {
    34  			typeName := strings.TrimPrefix(sc.exprString(spec.Recv.List[0].Type), "*")
    35  			methods, _ := sc.Methods[typeName]
    36  			methods = append(methods, GetMethodMeta(spec))
    37  			if sc.Methods == nil {
    38  				sc.Methods = make(map[string][]MethodMeta)
    39  			}
    40  			sc.Methods[typeName] = methods
    41  		}
    42  	case *ast.GenDecl:
    43  		if spec.Tok == token.CONST {
    44  			var comments []string
    45  			if spec.Doc != nil {
    46  				for _, comment := range spec.Doc.List {
    47  					comments = append(comments, strings.TrimSpace(strings.TrimPrefix(comment.Text, "//")))
    48  				}
    49  			}
    50  			var typeName string
    51  			for _, item := range spec.Specs {
    52  				valueSpec := item.(*ast.ValueSpec)
    53  				if len(valueSpec.Names) == 0 {
    54  					continue
    55  				}
    56  				switch specType := valueSpec.Type.(type) {
    57  				case *ast.Ident:
    58  					typeName = specType.Name
    59  				case nil:
    60  					// useless
    61  					//if len(valueSpec.Values) > 0 {
    62  					//	switch valueExpr := valueSpec.Values[0].(type) {
    63  					//	case *ast.BasicLit:
    64  					//		switch valueExpr.Kind {
    65  					//		case token.INT:
    66  					//			typeName = "int"
    67  					//		case token.FLOAT:
    68  					//			typeName = "float64"
    69  					//		case token.IMAG:
    70  					//			typeName = "complex128"
    71  					//		case token.CHAR:
    72  					//			typeName = "rune"
    73  					//		case token.STRING:
    74  					//			typeName = "string"
    75  					//		default:
    76  					//			continue
    77  					//		}
    78  					//	}
    79  					//}
    80  				}
    81  				sc.Consts[typeName] = append(sc.Consts[typeName], valueSpec.Names[0].Name)
    82  			}
    83  		}
    84  	}
    85  	return nil
    86  }
    87  
    88  // NewEnumCollector initializes an EnumCollector
    89  func NewEnumCollector(exprString func(ast.Expr) string) *EnumCollector {
    90  	return &EnumCollector{
    91  		Methods:    make(map[string][]MethodMeta),
    92  		Package:    PackageMeta{},
    93  		exprString: exprString,
    94  		Consts:     make(map[string][]string),
    95  	}
    96  }
    97  
    98  func EnumsOf(file string, exprString func(ast.Expr) string) *EnumCollector {
    99  	fset := token.NewFileSet()
   100  	root, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
   101  	if err != nil {
   102  		panic(err)
   103  	}
   104  	sc := NewEnumCollector(exprString)
   105  	ast.Walk(sc, root)
   106  	return sc
   107  }