Go 命令行解析 flag 包之擴展新類型

上篇文章 說到，除布爾類型 Flag，flag 支持的還有整型（int、int6四、uint、uint64）、浮點型（float64）、字符串（string）和時長（duration）。

flag 內置支持能知足大部分的需求，但某些場景，須要自定義解析規則。一個優秀的庫確定要支持擴展的。本文將介紹如何爲 flag 擴展一個新的類型支持？

擴展目標

在 gvg 這個小工具中，list 子命令支持獲取 Go 的版本列表。但版本的信息來源有多處，好比 installed（已安裝）、local（本地倉庫）和 remote（遠程倉庫）。

查看下 list 的幫助信息，以下：

NAME:
   gvg list - list go versions

USAGE:
   gvg list [command options] [arguments...]

OPTIONS:
   --origin value  the origin of version information , such as installed, local, remote (default: "installed")
複製代碼

能夠看出，list 子命令支持一個 Flag 選項，--origin。它用於指定版本信息的來源，容許值的範圍是 installed、local 和 remote。

若是要求不嚴格，用 StringVar 也能夠實現。但問題是，使用 String，即便輸入不在指定範圍也能成功解析，不夠嚴謹。雖然說在獲取後也能夠檢查，但仍是不夠靈活、可配置型也差。

接下來，咱們要實現一個新的類型的 Flag，使選項的值必需在指定範圍，不然要給出必定的錯誤提示信息。

實現思路

如何展一個新類型呢？

能夠參考 flag 包內置類型的實現思路，好比 flag.DurationVar。Duration 不是基礎類型，解析結果是存放到了 time.Duration 類型中，可能更有參考價值。

進入到 flag.DurationVar 查看源碼，以下：

func DurationVar(p *time.Duration, name string, value time.Duration, usage string) {
	CommandLine.Var(newDurationValue(value, p), name, usage)
}
複製代碼

經過 newDurationValue 建立了一個類型爲 durationValue 的變量，並傳入到了 CommandLine.Var 方法中。

若是繼續往下追，會根據 Value 建立一個 Flag 變量。 以下：

func (f *FlagSet) Var(value Value, name string, usage string) {
	flag := &Flag{name, usage, value, value.String()}
	...
}
複製代碼

從 Var 的定義能夠看出，它的第一個參數類型是 Value 接口類型，也就說，durationValue 是實現了 Value 接口的類型。

注意，源碼中出現的 FlagSet 能夠先忽略，它是下篇介紹子命令時重點關注的對象。

看下 Value 的定義，以下：

type Value interface {
	String() string
	Set(string) error
}
複製代碼

那麼，durationValue 的實現代碼如何？

// 傳入參數分別是默認值和獲取 Flag 值的變量地址
func newDurationValue(val time.Duration, p *time.Duration) *durationValue {
	// 將默認值設置到 p 上
	*p = val
	// 使用 p 建立新的類型，保證能夠獲取到解析的結果
	return (*durationValue)(p)
}

// Set 方法負責解析傳入的值
func (d *durationValue) Set(s string) error {
	v, err := time.ParseDuration(s)
	if err != nil {
		err = errParse
	}
	*d = durationValue(v)
	return err
}

// 獲取真正的值
func (d *durationValue) String() string { return (*time.Duration)(d).String() }
複製代碼

核心在兩個地方。

一個是建立新類型變量時，要使用傳入的變量地址建立新類型變量，以實現將解析結果放到其中，讓前端能獲取到，二是 Set 方法中實現命令行傳入字符串的解析。

邏輯梳理

看完上個小節，基本已經瞭解如何擴展一個新類型了。本質是是實現 Value 接口。

再看下以前提到的幾個變量，分別是存放解析結果的指針、解析命令行輸入的 Value 和表示一個選項的 Flag。對應於 flag.DurationVar，這個變量的類型分別是 *time.Duration、durationValue 和 Flag。

好比有 duration=1h，大體流程是首先從 os.Args 獲取參數，按規則解析出選項名稱 duration，查找是否存在名稱爲 duration 的 Flag，若是存在，使用 Flag.Value.Set 解析 1h，若是不知足 duration 的要求，將給出錯誤提示。

實現新類型

如今實現文章開頭要求的目標。

新類型定義以下：

type stringEnumValue struct {
	options []string
	p   *string
}
複製代碼

名爲 StringEnumValue，即字符串枚舉。它有 options 和 p 兩個成員，options 指定必定範圍的值，p 是 string 指針，保存解析結果的變量的地址。

下面定義建立 StringEnumValue 變量的函數 newStringEnumValue，代碼以下：

func newStringEnumValue(val string, p *string, options []string) *StringEnumValue {
	*option = val
	return &stringEnumValue{options: options, p: p}
}
複製代碼

除了 val 和 p 兩個必要的輸入外，還有一個 string 切片類型的數，名爲 options，它用於範圍的限定。而函數主體，首先設置默認值，而後使用 options 和 p 建立變量返回。

Set 是核心方法，解析命令行傳入字符串。代碼以下：

func (s *StringEnumValue) Set(v string) error {
	for _, option := range s.options {
		if v == option {
			*(s.p) = v
			return nil
		}
	}

	return fmt.Errorf("must be one of %v", s.options)
}
複製代碼

循環檢查輸入參數 v 是否知足要求。定義以下：

最後是 String() 方法，

func (s *StringEnumValue) String() string {
	return *(s.p)
}
複製代碼

返回 p 指針中的值。前面分析實現思路時，Flag 在設置默認值時就調用了它。

使用 StringEnumValue

直接看代碼吧。以下：

var origin string
func init() {
	flag.Var(
		newStringEnumValue(
			"installed", 	// 默認值
			&origin,
			[]string{"installed", "local", "remote"},
		),
		"origin",
		`the origin of version information, such as installed, local, remote (default: "installed")`,
	)
}

func main() {
	flag.Parse()
	fmt.Println(option)
}
複製代碼

重點就是 flag.Var(newStringEnumValue(...)，...)。若是以爲有點囉嗦，但願和其餘類型新建過程相同，在這個基礎上能夠再包裝。代碼以下：

func StringEnumVar(p *string, name string, options []string, defVal string, usage string) {
	flag.Var(newStringEnumValue(defVal, p, options), name, usage)
}
複製代碼

編譯測試下，結果以下：

$ gvg --origin=any
invalid value "any" for flag -origin: must be one of [installed local remote]
Usage of gvg:
  -origin value
  	the origin of version information, such as installed, local, remote (default installed)
$ gvg --origin=remote
origin remote
複製代碼

總結

本文介紹瞭如何爲 flag 擴展一個類型支持，經過分析源碼理清實現思路。最後建立了一個只接收指定範圍值的 Value 類型。

---

歡迎關注個人微信公衆號。