github.com/unionj-cloud/go-doudou@v1.3.8-0.20221011095552-0088008e5b31/cmd/internal/astutils/enumcollector.go (about)

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