Go flag 詳解,實現二級子命令

微信圖片_20210216141958

前言

日常開發使用到的命令行工具大都支援如下特性:

  • 文檔自動生成(如 -h --help
  • 多級子命令(如 docker exec -it
  • 支援參數(如 ls -color=auto
  • 長短選項(如 -v--verbose
  • 全局選項(如 docker -D run -d nginx
  • Tab 自動補全

本文就探討一下 Go 語言中如何寫一個擁有類似特性的命令行程式。

os.Args

類似於 shell 中的 $1 $2 ,在 Go 中可以使用 os.Args 來獲取命令行參數,這種臨時使用一兩個參數還可以,程式碼可維護性太差了,不推薦使用。

其中 Args[0] 是程式的名稱,Args[1] 是第一個參數,依此類推。

image-20220627173544369

flag 快速開始

Go 標準庫自帶的 flag 包可以實現簡單的命令行解析,我們模仿一下 ls 命令的參數,示例如下:

func main() {
	// 直接定義 flag,返回值為指針
	all := flag.Bool("all", true, "do not ignore entries starting with .")
	color := flag.String("color", "omitted", "colorize the output")

	// 也可以將 flag 綁定到變數
	var almostAll bool
	flag.BoolVar(&almostAll, "almost-all", false, "do not list implied . and ..")

	// 除過上面的 Bool String 還有 Int Float64 等其他常用類型

	flag.Parse()

	// Parse 後就可以獲取到具體參數的值
	fmt.Println(*all, *color, almostAll)
}

image-20220627175129721

可以看到非常簡單的幾行程式碼,就實現了一個還不錯的命令行小工具,支援 --- (效果是一致的),對於非 bool 類型的 flag 其值支援 -flag=val 或者 -flag val

長短選項

短選項書寫快捷,適合在終端下面執行,而長選項可讀性高,適合在腳本中書寫,通過共享一個變數,即可達到此效果

func main() {
	var name string
	defaultVal := "tom"
	useage := "your name"
	flag.StringVar(&name, "n", defaultVal, useage+" (shorthand)")
	flag.StringVar(&name, "name", defaultVal, useage)
	flag.Parse()

	fmt.Println(name)
}

image-20220627183431251

自定義類型

flag 也支援我們自定義參數的類型,方便我們對參數的格式,輸出形式做更加自由的處理,更好的封裝。

type Durations []time.Duration

func (d *Durations) String() string {
	return fmt.Sprint(*d)
}

func (d *Durations) Set(value string) error {
	// 支援逗號分割的參數,如:-d 1m,2s,1h
	// 也支援 -d 1m -d 2s -d 1h 這種寫法
	// 如果不想 -d 被指定多次,可以加上這段 if 邏輯
	// if len(*d) > 0 {
	// 	return errors.New("-d flag already set")
	// }

	for _, v := range strings.Split(value, ",") {
		duration, err := time.ParseDuration(v)
		if err != nil {
			return err
		}
		*d = append(*d, duration)
	}

	return nil
}

func main() {
	var param Durations
	// 第一個參數是介面類型,我們自定義的 Durations 只需要實現 String() 和 Set() 方法即可
	flag.Var(&param, "d", "time duration, comma-separated list")
	flag.Parse()
	fmt.Println(param)
}

image-20220628124237973

二級子命令

在說二級子命令前,我們先看一下 flag 的核心流程,幫助我們更好的理解二級子命令的實現。

// src/flag/flag.go

// 代表了每一個 flag,如 --name=tom
type Flag struct {
	Name     string // name as it appears on command line
	Usage    string // help message
	Value    Value  // value as set
	DefValue string // default value (as text); for usage message
}

// 代表本次命令輸出的所有 flag ,如 -l --size=10 --verbose
type FlagSet struct {
	Usage func()

	name          string
	parsed        bool
	actual        map[string]*Flag
	formal        map[string]*Flag
	args          []string // arguments after flags
	errorHandling ErrorHandling
	output        io.Writer // nil means stderr; use Output() accessor
}

在快速開始的程式碼中,核心程式碼就兩句,就對應了 flag 的流程,先註冊後解析。

color := flag.String("color", "omitted", "colorize the output")
flag.Parse()

追著 flag.String 可以看到它其實調用的是 CommandLine.StringCommandLine 是一個全局的 FlagSet 實例,最終 flag.String 會調用 FlagSetVar 方法,完成所有命令的註冊。

// src/flag/flag.go

func String(name string, value string, usage string) *string {
	return CommandLine.String(name, value, usage)
}

var CommandLine = NewFlagSet(os.Args[0], ExitOnError)

// 省略到校驗的一部分邏輯,可以看到核心就是 f.formal[name] = flag
func (f *FlagSet) Var(value Value, name string, usage string) {
	// ...

	// Remember the default value as a string; it won't change.
	flag := &Flag{name, usage, value, value.String()}
	// ...
    
	if f.formal == nil {
		f.formal = make(map[string]*Flag)
	}
	f.formal[name] = flag
}

flag.Parse 則最終調用的是 FlagSetParse 方法,完成實際輸入值的解析。

func (f *FlagSet) Parse(arguments []string) error {
	f.parsed = true
	f.args = arguments
	for { // 循環直到所有的 flag 都解析完成
		seen, err := f.parseOne()
		if seen {
			continue
		}
		// ...
	}
	return nil
}

在理解了上面了流程後,我們就基於 FlagSet 來模仿一下 docker 的二級命令,程式碼如下:

type MyFlagSet struct {
	*flag.FlagSet
	cmdComment string // 二級子命令本身的注釋
}

func main() {
	// docker ps
	psCmd := &MyFlagSet{
		FlagSet:    flag.NewFlagSet("ps", flag.ExitOnError),
		cmdComment: "List containers",
	}
	psCmd.Bool("a", false, "Show all containers (default shows just running)")
	psCmd.Bool("s", false, "Display total file sizes")

	// docker run
	runCmd := &MyFlagSet{
		FlagSet:    flag.NewFlagSet("run", flag.ExitOnError),
		cmdComment: "Run a command in a new container",
	}
	runCmd.Int("c", 1, "CPU shares (relative weight)")
	runCmd.String("name", "", "Assign a name to the container")

	// 用 map 保存所有的二級子命令,方便快速查找
	subcommands := map[string]*MyFlagSet{
		psCmd.Name():  psCmd,
		runCmd.Name(): runCmd,
	}

	useage := func() { // 整個命令行的幫助資訊
		fmt.Printf("Usage: docker COMMAND\n\n")
		for _, v := range subcommands {
			fmt.Printf("%s %s\n", v.Name(), v.cmdComment)
			v.PrintDefaults() // 使用 flag 庫自帶的格式輸出子命令的選項幫助資訊
			fmt.Println()
		}
		os.Exit(2)
	}

	if len(os.Args) < 2 { // 即沒有輸入子命令
		useage()
	}

	cmd := subcommands[os.Args[1]] // 第二個參數必須是我們支援的子命令
	if cmd == nil {
		useage()
	}

    cmd.Parse(os.Args[2:]) // 注意這裡是 cmd.Parse 不是 flag.Parse,且值是 Args[2:]

	// 輸出解析後的結果
	fmt.Println("command name is:", cmd.Name())
	cmd.Visit(func(f *flag.Flag) {
		fmt.Printf("option %s, value is %s\n", f.Name, f.Value)
	})
}

image-20220629083323738

可以看到效果還不錯呢,到目前為止,除了全局選項和自動補全未實現,其他的特性都有了,可以看到總體來說 flag 簡單易用,對於一些小程式來說完全足夠了。

當然對於比較複雜的程式,還是推薦使用更加強大的 cobra,可以參考筆者寫的 Markdown 小幫手 marker

參考

//pkg.go.dev/flag

//gobyexample.com/command-line-subcommands