github.com/15mga/kiwi@v0.0.2-0.20240324021231-b95d5c3ac751/var.go (about)

     1  package kiwi
     2  
     3  import (
     4  	"flag"
     5  	"fmt"
     6  	"github.com/15mga/kiwi/util"
     7  	"os"
     8  	"strconv"
     9  	"strings"
    10  
    11  	"github.com/15mga/kiwi/ds"
    12  )
    13  
    14  type varItem struct {
    15  	name  string
    16  	usage string
    17  	val   any
    18  }
    19  
    20  type Var interface {
    21  	int | int64 | float64 | bool | string
    22  }
    23  
    24  var (
    25  	_VarMap = ds.NewKSet[string, *varItem](4, func(item *varItem) string {
    26  		return item.name
    27  	})
    28  )
    29  
    30  func AddVar[T Var](name string, def T, usage string) {
    31  	_VarMap.Set(&varItem{
    32  		name:  name,
    33  		val:   def,
    34  		usage: usage,
    35  	})
    36  }
    37  
    38  func ParseVar() {
    39  	parseFlag()
    40  	parseEnv()
    41  	m := util.M{}
    42  	for _, v := range _VarMap.Values() {
    43  		m[v.name] = v.val
    44  	}
    45  }
    46  
    47  func parseFlag() {
    48  	m := make(map[string]any)
    49  	for _, item := range _VarMap.Values() {
    50  		switch d := item.val.(type) {
    51  		case int:
    52  			var v int
    53  			flag.IntVar(&v, item.name, d, item.usage)
    54  			m[item.name] = &v
    55  		case int64:
    56  			var v int64
    57  			flag.Int64Var(&v, item.name, d, item.usage)
    58  			m[item.name] = &v
    59  		case float64:
    60  			var v float64
    61  			flag.Float64Var(&v, item.name, d, item.usage)
    62  			m[item.name] = &v
    63  		case bool:
    64  			var v bool
    65  			flag.BoolVar(&v, item.name, d, item.usage)
    66  			m[item.name] = &v
    67  		case string:
    68  			var v string
    69  			flag.StringVar(&v, item.name, d, item.usage)
    70  			m[item.name] = &v
    71  		}
    72  	}
    73  	flag.Parse()
    74  	for _, item := range _VarMap.Values() {
    75  		v := m[item.name]
    76  		switch d := v.(type) {
    77  		case *int:
    78  			item.val = *d
    79  		case *int64:
    80  			item.val = *d
    81  		case *float64:
    82  			item.val = *d
    83  		case *bool:
    84  			item.val = *d
    85  		case *string:
    86  			item.val = *d
    87  		}
    88  	}
    89  }
    90  
    91  func parseEnv() {
    92  	for _, item := range _VarMap.Values() {
    93  		v, ok := os.LookupEnv(strings.ToUpper(item.name))
    94  		if ok {
    95  			switch item.val.(type) {
    96  			case int:
    97  				i, err := strconv.Atoi(v)
    98  				if err != nil {
    99  					fmt.Println(err.Error())
   100  				}
   101  				item.val = i
   102  			case int64:
   103  				i, err := strconv.ParseInt(v, 10, 64)
   104  				if err != nil {
   105  					fmt.Println(err.Error())
   106  				}
   107  				item.val = i
   108  			case float64:
   109  				i, err := strconv.ParseFloat(v, 64)
   110  				if err != nil {
   111  					fmt.Println(err.Error())
   112  				}
   113  				item.val = i
   114  			case bool:
   115  				item.val = strings.ToLower(v) == "true"
   116  			case string:
   117  				item.val = v
   118  			}
   119  		}
   120  	}
   121  }
   122  
   123  func GetVar[T Var](name string) (T, bool) {
   124  	o, ok := _VarMap.Get(name)
   125  	if !ok {
   126  		return util.Default[T](), false
   127  	}
   128  	v, ok := o.val.(T)
   129  	return v, ok
   130  }