flagset.go 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273
  1. // mauflag - An extendable command-line argument parser for Golang
  2. // Copyright (C) 2016 Maunium
  3. //
  4. // This program is free software: you can redistribute it and/or modify
  5. // it under the terms of the GNU General Public License as published by
  6. // the Free Software Foundation, either version 3 of the License, or
  7. // (at your option) any later version.
  8. // This program is distributed in the hope that it will be useful,
  9. // but WITHOUT ANY WARRANTY; without even the implied warranty of
  10. // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  11. // GNU General Public License for more details.
  12. // You should have received a copy of the GNU General Public License
  13. // along with this program. If not, see <http://www.gnu.org/licenses/>.
  14. package mauflag
  15. import (
  16. "fmt"
  17. "os"
  18. "strings"
  19. )
  20. // Set is a set of flags with certain input arguments
  21. type Set struct {
  22. // The list of strings used as input
  23. InputArgs []string
  24. // Whether or not to ignore all flags after the user has entered two dashes with no flag key ("--")
  25. // If enabled, all arguments after two dashes with no flag key will go into the args array (@see Args())
  26. DoubleLineEscape bool
  27. // Whether or not to exit the program when there's an error
  28. // If enabled, the error message will be printed to `stderr` after which `os.Exit(1)` will be called.
  29. ExitOnError bool
  30. wantHelp *bool
  31. basicUsage string
  32. helpFirstLine string
  33. args []string
  34. flags []*Flag
  35. }
  36. // New creates a new flagset
  37. func New(args []string) *Set {
  38. return &Set{InputArgs: args, DoubleLineEscape: true, ExitOnError: false}
  39. }
  40. // Args returns the arguments that weren't associated with any flag
  41. func (fs *Set) Args() []string {
  42. return fs.args
  43. }
  44. // Arg returns the string at the given index from the list Args() returns
  45. // If the index does not exist, Arg will return an empty string.
  46. func (fs *Set) Arg(i int) string {
  47. if len(fs.args) <= i {
  48. return ""
  49. }
  50. return fs.args[i]
  51. }
  52. // NArg returns the number of arguments not associated with any flags
  53. func (fs *Set) NArg() int {
  54. return len(fs.args)
  55. }
  56. // MakeHelpFlag creates the -h, --help flag
  57. func (fs *Set) MakeHelpFlag() (*bool, *Flag) {
  58. var flag = fs.Make().Key("h", "help").Usage("Show this help page.").UsageCategory("Help")
  59. fs.wantHelp = flag.Bool()
  60. return fs.wantHelp, flag
  61. }
  62. // CheckHelpFlag checks if the help flag is set and prints the help page if needed.
  63. // Return value tells whether or not the help page was printed
  64. func (fs *Set) CheckHelpFlag() bool {
  65. if fs.wantHelp != nil && *fs.wantHelp {
  66. fs.PrintHelp()
  67. return true
  68. }
  69. return false
  70. }
  71. // SetHelpTitles sets the first line (program name and basic explanation) and basic usage specification
  72. func (fs *Set) SetHelpTitles(firstLine, basicUsage string) {
  73. fs.helpFirstLine = firstLine
  74. fs.basicUsage = basicUsage
  75. }
  76. // PrintHelp prints the help page
  77. func (fs *Set) PrintHelp() {
  78. var helpSectNames = make(map[string]int)
  79. var helpSectIndexes = make(map[int]string)
  80. var helpSects = make([][]*Flag, 0)
  81. for _, flag := range fs.flags {
  82. index, ok := helpSectNames[flag.usageCat]
  83. if !ok {
  84. arr := []*Flag{flag}
  85. helpSects = append(helpSects, arr)
  86. helpSectNames[flag.usageCat] = len(helpSects) - 1
  87. helpSectIndexes[len(helpSects)-1] = flag.usageCat
  88. } else {
  89. helpSects[index] = append(helpSects[index], flag)
  90. }
  91. }
  92. data, maxLen := fs.formatFlagHelp(helpSects)
  93. fmt.Printf(`%s
  94. Usage:
  95. %s
  96. `, fs.helpFirstLine, fs.basicUsage)
  97. for sect, sData := range data {
  98. fmt.Print(helpSectIndexes[sect], " options:\n")
  99. for i, fData := range sData {
  100. fmt.Print(fData, strings.Repeat(" ", maxLen-len(fData)+3), helpSects[sect][i].usage, "\n")
  101. }
  102. fmt.Print("\n")
  103. }
  104. }
  105. func (fs *Set) formatFlagHelp(helpSects [][]*Flag) (data [][]string, maxLen int) {
  106. maxLen = 0
  107. data = make([][]string, len(helpSects))
  108. for sect, flags := range helpSects {
  109. var sData = make([]string, len(flags))
  110. for i, flag := range flags {
  111. var fData = []string{" "}
  112. for _, key := range flag.shortKeys {
  113. fData = append(fData, "-")
  114. fData = append(fData, key)
  115. fData = append(fData, ", ")
  116. }
  117. for _, key := range flag.longKeys {
  118. fData = append(fData, "--")
  119. fData = append(fData, key)
  120. if len(flag.usageValName) > 0 {
  121. fData = append(fData, "=")
  122. fData = append(fData, flag.usageValName)
  123. }
  124. fData = append(fData, ", ")
  125. }
  126. if fData[len(fData)-1] == ", " {
  127. fData = fData[:len(fData)-1]
  128. }
  129. sData[i] = strings.Join(fData, "")
  130. if len(sData[i]) > maxLen {
  131. maxLen = len(sData[i])
  132. }
  133. }
  134. data[sect] = sData
  135. }
  136. return
  137. }
  138. func (fs *Set) err(format string, args ...interface{}) error {
  139. if fs.ExitOnError {
  140. fmt.Fprintf(os.Stderr, format, args...)
  141. fmt.Fprint(os.Stderr, "\n")
  142. os.Exit(1)
  143. return nil
  144. }
  145. return fmt.Errorf(format, args...)
  146. }
  147. // Parse the input arguments in this flagset into mauflag form
  148. // Before this function is called all the flags will have either the type default or the given default value
  149. func (fs *Set) Parse() error {
  150. var flag *Flag
  151. var key string
  152. var noMoreFlags = false
  153. for _, arg := range fs.InputArgs {
  154. if noMoreFlags {
  155. fs.args = append(fs.args, arg)
  156. } else if arg == "--" && fs.DoubleLineEscape {
  157. noMoreFlags = true
  158. } else if flag != nil {
  159. err := flag.setValue(arg)
  160. if err != nil {
  161. return fs.err("Flag %s was not a %s", key, flag.Value.Name())
  162. }
  163. flag = nil
  164. } else if arg[0] == '-' && len(arg) > 1 {
  165. arg = strings.ToLower(arg)
  166. var err error
  167. flag, key, err = fs.flagFound(arg[1:])
  168. if err != nil {
  169. return err
  170. }
  171. } else {
  172. fs.args = append(fs.args, arg)
  173. }
  174. }
  175. return nil
  176. }
  177. func (fs *Set) flagFound(arg string) (flag *Flag, key string, err error) {
  178. key = arg
  179. var val string
  180. if strings.ContainsRune(key, '=') {
  181. val = key[strings.Index(key, "=")+1:]
  182. key = key[:strings.Index(key, "=")]
  183. }
  184. var rem string
  185. flag, key, rem = fs.getFlag(key)
  186. flag, key, rem, err = fs.handleFlag(flag, key, rem, val)
  187. if len(rem) > 0 {
  188. return fs.flagFound(rem)
  189. }
  190. return
  191. }
  192. func (fs *Set) handleFlag(flag *Flag, key, rem, val string) (*Flag, string, string, error) {
  193. var err error
  194. var isBool bool
  195. if flag != nil {
  196. _, isBool = flag.Value.(*boolValue)
  197. }
  198. if flag == nil {
  199. err = fs.err("Unknown flag: %s", key)
  200. } else if len(val) > 0 && len(rem) == 0 {
  201. flag.setValue(val)
  202. } else if isBool {
  203. flag.setValue("true")
  204. } else if len(rem) > 0 {
  205. flag.setValue(rem)
  206. rem = ""
  207. } else {
  208. return flag, key, rem, err
  209. }
  210. return nil, "", rem, err
  211. }
  212. func (fs *Set) getFlag(key string) (*Flag, string, string) {
  213. if key[0] == '-' {
  214. key = key[1:]
  215. for _, lflag := range fs.flags {
  216. for _, lkey := range lflag.longKeys {
  217. if lkey == key {
  218. return lflag, lkey, ""
  219. }
  220. }
  221. }
  222. return nil, key, ""
  223. }
  224. rem := key[1:]
  225. key = key[0:1]
  226. for _, lflag := range fs.flags {
  227. for _, lkey := range lflag.shortKeys {
  228. if lkey == key {
  229. return lflag, lkey, rem
  230. }
  231. }
  232. }
  233. return nil, key, rem
  234. }