string_to_int.go 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  1. package pflag
  2. import (
  3. "bytes"
  4. "fmt"
  5. "strconv"
  6. "strings"
  7. )
  8. // -- stringToInt Value
  9. type stringToIntValue struct {
  10. value *map[string]int
  11. changed bool
  12. }
  13. func newStringToIntValue(val map[string]int, p *map[string]int) *stringToIntValue {
  14. ssv := new(stringToIntValue)
  15. ssv.value = p
  16. *ssv.value = val
  17. return ssv
  18. }
  19. // Format: a=1,b=2
  20. func (s *stringToIntValue) Set(val string) error {
  21. ss := strings.Split(val, ",")
  22. out := make(map[string]int, len(ss))
  23. for _, pair := range ss {
  24. kv := strings.SplitN(pair, "=", 2)
  25. if len(kv) != 2 {
  26. return fmt.Errorf("%s must be formatted as key=value", pair)
  27. }
  28. var err error
  29. out[kv[0]], err = strconv.Atoi(kv[1])
  30. if err != nil {
  31. return err
  32. }
  33. }
  34. if !s.changed {
  35. *s.value = out
  36. } else {
  37. for k, v := range out {
  38. (*s.value)[k] = v
  39. }
  40. }
  41. s.changed = true
  42. return nil
  43. }
  44. func (s *stringToIntValue) Type() string {
  45. return "stringToInt"
  46. }
  47. func (s *stringToIntValue) String() string {
  48. var buf bytes.Buffer
  49. i := 0
  50. for k, v := range *s.value {
  51. if i > 0 {
  52. buf.WriteRune(',')
  53. }
  54. buf.WriteString(k)
  55. buf.WriteRune('=')
  56. buf.WriteString(strconv.Itoa(v))
  57. i++
  58. }
  59. return "[" + buf.String() + "]"
  60. }
  61. func stringToIntConv(val string) (interface{}, error) {
  62. val = strings.Trim(val, "[]")
  63. // An empty string would cause an empty map
  64. if len(val) == 0 {
  65. return map[string]int{}, nil
  66. }
  67. ss := strings.Split(val, ",")
  68. out := make(map[string]int, len(ss))
  69. for _, pair := range ss {
  70. kv := strings.SplitN(pair, "=", 2)
  71. if len(kv) != 2 {
  72. return nil, fmt.Errorf("%s must be formatted as key=value", pair)
  73. }
  74. var err error
  75. out[kv[0]], err = strconv.Atoi(kv[1])
  76. if err != nil {
  77. return nil, err
  78. }
  79. }
  80. return out, nil
  81. }
  82. // GetStringToInt return the map[string]int value of a flag with the given name
  83. func (f *FlagSet) GetStringToInt(name string) (map[string]int, error) {
  84. val, err := f.getFlagType(name, "stringToInt", stringToIntConv)
  85. if err != nil {
  86. return map[string]int{}, err
  87. }
  88. return val.(map[string]int), nil
  89. }
  90. // StringToIntVar defines a string flag with specified name, default value, and usage string.
  91. // The argument p points to a map[string]int variable in which to store the values of the multiple flags.
  92. // The value of each argument will not try to be separated by comma
  93. func (f *FlagSet) StringToIntVar(p *map[string]int, name string, value map[string]int, usage string) {
  94. f.VarP(newStringToIntValue(value, p), name, "", usage)
  95. }
  96. // StringToIntVarP is like StringToIntVar, but accepts a shorthand letter that can be used after a single dash.
  97. func (f *FlagSet) StringToIntVarP(p *map[string]int, name, shorthand string, value map[string]int, usage string) {
  98. f.VarP(newStringToIntValue(value, p), name, shorthand, usage)
  99. }
  100. // StringToIntVar defines a string flag with specified name, default value, and usage string.
  101. // The argument p points to a map[string]int variable in which to store the value of the flag.
  102. // The value of each argument will not try to be separated by comma
  103. func StringToIntVar(p *map[string]int, name string, value map[string]int, usage string) {
  104. CommandLine.VarP(newStringToIntValue(value, p), name, "", usage)
  105. }
  106. // StringToIntVarP is like StringToIntVar, but accepts a shorthand letter that can be used after a single dash.
  107. func StringToIntVarP(p *map[string]int, name, shorthand string, value map[string]int, usage string) {
  108. CommandLine.VarP(newStringToIntValue(value, p), name, shorthand, usage)
  109. }
  110. // StringToInt defines a string flag with specified name, default value, and usage string.
  111. // The return value is the address of a map[string]int variable that stores the value of the flag.
  112. // The value of each argument will not try to be separated by comma
  113. func (f *FlagSet) StringToInt(name string, value map[string]int, usage string) *map[string]int {
  114. p := map[string]int{}
  115. f.StringToIntVarP(&p, name, "", value, usage)
  116. return &p
  117. }
  118. // StringToIntP is like StringToInt, but accepts a shorthand letter that can be used after a single dash.
  119. func (f *FlagSet) StringToIntP(name, shorthand string, value map[string]int, usage string) *map[string]int {
  120. p := map[string]int{}
  121. f.StringToIntVarP(&p, name, shorthand, value, usage)
  122. return &p
  123. }
  124. // StringToInt defines a string flag with specified name, default value, and usage string.
  125. // The return value is the address of a map[string]int variable that stores the value of the flag.
  126. // The value of each argument will not try to be separated by comma
  127. func StringToInt(name string, value map[string]int, usage string) *map[string]int {
  128. return CommandLine.StringToIntP(name, "", value, usage)
  129. }
  130. // StringToIntP is like StringToInt, but accepts a shorthand letter that can be used after a single dash.
  131. func StringToIntP(name, shorthand string, value map[string]int, usage string) *map[string]int {
  132. return CommandLine.StringToIntP(name, shorthand, value, usage)
  133. }