一个命令变多后,main 很容易乱
Go 很适合写命令行工具。最开始你可能只有一个参数:
todo -file tasks.json
后来需求变成:
todo add -title "学习 Go"
todo list
todo done -id 1
这时如果还把所有参数都塞进全局 flag,main.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 需要 title,list 不需要。这样帮助信息也更准确。
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)
}
}
把 args、stdout、stderr 都作为参数传入,是命令行工具可测试的关键。
帮助信息也要可维护
CLI 工具做久了,帮助信息会变得很重要。用户不会总是打开 README,他们更可能先运行 tool help 或 tool 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。入门阶段先用标准库写清楚结构,会更理解命令行工具的本质。
继续阅读
探索更多技术文章
浏览归档,发现更多关于系统设计、工具链和工程实践的内容。