Go CLI 子命令入门:用 flag.FlagSet 写清楚的命令行工具

本文讲解如何用 Go 标准库 flag.FlagSet 实现 add、list、done 这类子命令,帮助初学者组织可维护命令行工具。

一个命令变多后,main 很容易乱

Go 很适合写命令行工具。最开始你可能只有一个参数:

todo -file tasks.json

后来需求变成:

todo add -title "学习 Go"
todo list
todo done -id 1

这时如果还把所有参数都塞进全局 flagmain.go 会很快变乱。标准库 flag.FlagSet 可以为每个子命令单独解析参数。你不一定要马上引入 Cobra 这类库,小工具用标准库就能组织得很清楚。

这篇文章写一个 todo CLI 的子命令骨架。

run 函数接收 args

main 保持薄:

func main() {
	if err := run(os.Args[1:], os.Stdout, os.Stderr); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

run 分发子命令:

func run(args []string, stdout io.Writer, stderr io.Writer) error {
	if len(args) == 0 {
		return fmt.Errorf("command is required: add, list, done")
	}

	switch args[0] {
	case "add":
		return runAdd(args[1:], stdout, stderr)
	case "list":
		return runList(args[1:], stdout, stderr)
	case "done":
		return runDone(args[1:], stdout, stderr)
	default:
		return fmt.Errorf("unknown command: %s", args[0])
	}
}

这样测试时不需要真的启动进程,只要调用 run

add 子命令

func runAdd(args []string, stdout io.Writer, stderr io.Writer) error {
	fs := flag.NewFlagSet("add", flag.ContinueOnError)
	fs.SetOutput(stderr)

	title := fs.String("title", "", "task title")
	file := fs.String("file", "tasks.json", "task file")

	if err := fs.Parse(args); err != nil {
		return err
	}
	if strings.TrimSpace(*title) == "" {
		return fmt.Errorf("title is required")
	}

	task := Task{Title: strings.TrimSpace(*title)}
	if err := appendTask(*file, task); err != nil {
		return err
	}

	fmt.Fprintf(stdout, "added: %s\n", task.Title)
	return nil
}

flag.ContinueOnError 表示解析错误时返回 error,而不是直接退出程序。fs.SetOutput(stderr) 让帮助和错误输出进入传入的 stderr,测试更容易。

list 子命令

func runList(args []string, stdout io.Writer, stderr io.Writer) error {
	fs := flag.NewFlagSet("list", flag.ContinueOnError)
	fs.SetOutput(stderr)

	file := fs.String("file", "tasks.json", "task file")
	if err := fs.Parse(args); err != nil {
		return err
	}

	tasks, err := loadTasks(*file)
	if err != nil {
		return err
	}
	if len(tasks) == 0 {
		fmt.Fprintln(stdout, "no tasks")
		return nil
	}

	for i, task := range tasks {
		status := " "
		if task.Done {
			status = "x"
		}
		fmt.Fprintf(stdout, "[%s] %d %s\n", status, i+1, task.Title)
	}
	return nil
}

每个子命令只关心自己的参数。add 需要 titlelist 不需要。这样帮助信息也更准确。

done 子命令

func runDone(args []string, stdout io.Writer, stderr io.Writer) error {
	fs := flag.NewFlagSet("done", flag.ContinueOnError)
	fs.SetOutput(stderr)

	id := fs.Int("id", 0, "task id")
	file := fs.String("file", "tasks.json", "task file")
	if err := fs.Parse(args); err != nil {
		return err
	}
	if *id <= 0 {
		return fmt.Errorf("id is required")
	}

	if err := markDone(*file, *id); err != nil {
		return err
	}
	fmt.Fprintf(stdout, "done: %d\n", *id)
	return nil
}

参数校验放在子命令里,错误消息尽量直接。用户运行错命令时,应该知道该改什么。

测试子命令

func TestRunUnknownCommand(t *testing.T) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := run([]string{"bad"}, &stdout, &stderr)
	if err == nil {
		t.Fatal("expected error")
	}
}

测试 add 缺 title:

func TestRunAddMissingTitle(t *testing.T) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := run([]string{"add"}, &stdout, &stderr)
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "title") {
		t.Fatalf("error = %v", err)
	}
}

argsstdoutstderr 都作为参数传入,是命令行工具可测试的关键。

帮助信息也要可维护

CLI 工具做久了,帮助信息会变得很重要。用户不会总是打开 README,他们更可能先运行 tool helptool add -h。标准库 flag.FlagSet 可以给每个子命令单独设置输出位置和用法:

func newAddFlagSet(stderr io.Writer) (*flag.FlagSet, *string) {
	fs := flag.NewFlagSet("add", flag.ContinueOnError)
	fs.SetOutput(stderr)

	title := fs.String("title", "", "task title")
	fs.Usage = func() {
		fmt.Fprintln(stderr, "Usage: task add -title <title>")
		fs.PrintDefaults()
	}
	return fs, title
}

这样测试错误输出也很方便:

func TestAddHelp(t *testing.T) {
	var stderr bytes.Buffer
	fs, _ := newAddFlagSet(&stderr)
	fs.Usage()
	if !strings.Contains(stderr.String(), "task add") {
		t.Fatalf("help = %q", stderr.String())
	}
}

当命令行工具被脚本调用时,退出码也要稳定。参数错误一般返回非零,业务执行失败也返回非零,但成功且没有数据时是否算错误,要提前定义好。标准库不会替你设计这些规则,但它给了足够的结构让你把规则写清楚。

小结

Go 标准库的 flag.FlagSet 足够组织很多小型 CLI 子命令。main 负责退出码,run 负责分发,每个子命令有自己的 FlagSet 和校验逻辑。输入输出通过 io.Writer 注入,测试就不需要真实终端。

当命令很多、帮助复杂、需要自动补全时,可以再考虑 Cobra。入门阶段先用标准库写清楚结构,会更理解命令行工具的本质。

继续阅读

探索更多技术文章

浏览归档,发现更多关于系统设计、工具链和工程实践的内容。

全部文章 返回首页