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

     1  package enumgen
     2  
     3  import (
     4  	"go/ast"
     5  	"go/constant"
     6  	"go/types"
     7  	"sort"
     8  	"strconv"
     9  	"strings"
    10  
    11  	"github.com/machinefi/w3bstream/pkg/depends/x/pkgx"
    12  	"github.com/machinefi/w3bstream/pkg/depends/x/stringsx"
    13  )
    14  
    15  type Scanner struct {
    16  	pkg *pkgx.Pkg
    17  	res map[*types.TypeName]Options
    18  }
    19  
    20  func NewScanner(pkg *pkgx.Pkg) *Scanner { return &Scanner{pkg: pkg} }
    21  
    22  func (s *Scanner) Options(tn *types.TypeName) (Options, bool) {
    23  	if tn == nil {
    24  		return nil, false
    25  	}
    26  	if options, ok := s.res[tn]; ok {
    27  		return options, ok
    28  	}
    29  
    30  	pkg := s.pkg.PkgByPath(tn.Pkg().Path())
    31  	if pkg == nil {
    32  		return nil, false
    33  	}
    34  
    35  	for ident, def := range pkg.TypesInfo.Defs {
    36  		c, ok := def.(*types.Const)
    37  		if !ok {
    38  			continue
    39  		}
    40  		if c.Type() != tn.Type() {
    41  			continue
    42  		}
    43  		name := c.Name()
    44  		if strings.HasPrefix(name, "_") {
    45  			continue
    46  		}
    47  		spec := ident.Obj.Decl.(*ast.ValueSpec)
    48  		label := strings.TrimSpace(spec.Comment.Text())
    49  		val := c.Val()
    50  
    51  		switch val.Kind() {
    52  		case constant.String:
    53  			v, _ := strconv.Unquote(val.String())
    54  			s.Append(tn, NewStringOption(v, label))
    55  		case constant.Float:
    56  			v, _ := strconv.ParseFloat(val.String(), 64)
    57  			s.Append(tn, NewFloatOption(v, label))
    58  		case constant.Int:
    59  			// TYPE_NAME_UNKNOWN
    60  			// TYPE_NAME__XXX
    61  			if strings.HasPrefix(name, stringsx.UpperSnakeCase(tn.Name())) {
    62  				parts := strings.SplitN(name, "__", 2)
    63  				if len(parts) == 2 {
    64  					v, _ := strconv.ParseInt(val.String(), 10, 64)
    65  					s.Append(tn, NewOption(v, parts[1], label))
    66  				}
    67  			} else {
    68  				v, _ := strconv.ParseInt(val.String(), 10, 64)
    69  				s.Append(tn, NewIntOption(v, label))
    70  			}
    71  		default:
    72  			return nil, false
    73  		}
    74  	}
    75  	return s.res[tn], len(s.res[tn]) > 0
    76  }
    77  
    78  func (s *Scanner) Append(tn *types.TypeName, opt *Option) {
    79  	if s.res == nil {
    80  		s.res = make(map[*types.TypeName]Options)
    81  	}
    82  	s.res[tn] = append(s.res[tn], *opt)
    83  	sort.Sort(s.res[tn])
    84  }
    85  
    86  func (s *Scanner) Package() *pkgx.Pkg { return s.pkg }