github.com/octohelm/storage@v0.0.0-20240516030302-1ac2cc1ea347/devpkg/enumgen/enum_type.go (about)

     1  package enumgen
     2  
     3  import (
     4  	"fmt"
     5  	"go/constant"
     6  	"go/types"
     7  	"sort"
     8  	"strconv"
     9  	"strings"
    10  
    11  	"github.com/octohelm/gengo/pkg/gengo"
    12  )
    13  
    14  type EnumTypes map[string]map[types.Type]*EnumType
    15  
    16  func (e EnumTypes) ResolveEnumType(t types.Type) (*EnumType, bool) {
    17  	if n, ok := t.(*types.Named); ok {
    18  		if enumTypes, ok := e[n.Obj().Pkg().Path()]; ok && enumTypes != nil {
    19  			if enumType, ok := enumTypes[t]; ok {
    20  				return enumType, ok
    21  			}
    22  		}
    23  	}
    24  	return nil, false
    25  }
    26  
    27  func (e EnumTypes) Walk(gc gengo.Context, inPkgPath string) {
    28  	p := gc.Package(inPkgPath)
    29  
    30  	constants := p.Constants()
    31  
    32  	for k := range p.Constants() {
    33  		if e[inPkgPath] == nil {
    34  			e[inPkgPath] = map[types.Type]*EnumType{}
    35  		}
    36  
    37  		constv := constants[k]
    38  
    39  		if e[inPkgPath][constv.Type()] == nil {
    40  			e[inPkgPath][constv.Type()] = &EnumType{}
    41  		}
    42  
    43  		e[inPkgPath][constv.Type()].Add(constv, p.Comment(constv.Pos())...)
    44  	}
    45  }
    46  
    47  type EnumType struct {
    48  	ConstUnknown *types.Const
    49  	Constants    []*types.Const
    50  	Comments     map[*types.Const][]string
    51  }
    52  
    53  func (e *EnumType) IsIntStringer() bool {
    54  	return e.ConstUnknown != nil && len(e.Constants) > 0
    55  }
    56  
    57  func (e *EnumType) Label(cv *types.Const) string {
    58  	if comments, ok := e.Comments[cv]; ok {
    59  		label := strings.Join(comments, "")
    60  
    61  		if label != "" {
    62  			return label
    63  		}
    64  	}
    65  
    66  	return fmt.Sprintf("%v", e.Value(cv))
    67  }
    68  
    69  func (e *EnumType) Value(cv *types.Const) any {
    70  	if named, ok := cv.Type().(*types.Named); ok {
    71  		var parts = strings.SplitN(cv.Name(), "__", 2)
    72  
    73  		if len(parts) == 2 && parts[0] == gengo.UpperSnakeCase(named.Obj().Name()) {
    74  			return parts[1]
    75  		}
    76  	}
    77  
    78  	val := cv.Val()
    79  
    80  	if val.Kind() == constant.Int {
    81  		i, _ := strconv.ParseInt(val.String(), 10, 64)
    82  		return i
    83  	}
    84  
    85  	s, _ := strconv.Unquote(val.String())
    86  	return s
    87  }
    88  
    89  func (e *EnumType) Add(cv *types.Const, comments ...string) {
    90  	if e.Comments == nil {
    91  		e.Comments = map[*types.Const][]string{}
    92  	}
    93  
    94  	n := cv.Name()
    95  
    96  	if n[0] == '_' {
    97  		return
    98  	}
    99  
   100  	if strings.HasSuffix(n, "_UNKNOWN") {
   101  		e.ConstUnknown = cv
   102  		return
   103  	}
   104  
   105  	e.Comments[cv] = comments
   106  
   107  	parts := strings.SplitN(n, "__", 2)
   108  
   109  	if len(parts) == 2 {
   110  		names := strings.Split(cv.Type().String(), ".")
   111  		name := names[len(names)-1]
   112  
   113  		if gengo.UpperSnakeCase(name) == parts[0] {
   114  			e.Constants = append(e.Constants, cv)
   115  		}
   116  	} else {
   117  		e.Constants = append(e.Constants, cv)
   118  	}
   119  
   120  	sort.Slice(e.Constants, func(i, j int) bool {
   121  		return e.Constants[i].Val().String() < e.Constants[j].Val().String()
   122  	})
   123  }