GoのCobraを使ってCLIを作ってみた

はじめに

前回、flagパッケージを使ってCLIを作ってみたので、今回はCobraライブラリを使用してCLIを作ってみました。

コード

ルートのコマンドとして、toolkitという名前のコマンドを作成しています。
そのサブコマンドとして、countとlinesを作りました。
Useフィールドにコマンド名、Runフィールドに実行する関数を登録します。
そして、それぞれFlags().StringP()で引数を定義します。
定義できたら、最後にAddCommand()で登録します。

main.go

package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"

	"github.com/spf13/cobra"
)

func main() {
	var rootCmd = &cobra.Command{Use: "toolkit"}

	var countCmd = &cobra.Command{
		Use:   "count [flags]",
		Short: "Count occurrences of a word in a file",
		Run:   countWord,
	}

	var linesCmd = &cobra.Command{
		Use:   "lines [flags]",
		Short: "Count lines in a file",
		Run:   countLines,
	}

	countCmd.Flags().StringP("file", "f", "", "File to search in")
	countCmd.Flags().StringP("word", "w", "", "Word to search for")
	countCmd.MarkFlagRequired("file")
	countCmd.MarkFlagRequired("word")

	linesCmd.Flags().StringP("file", "f", "", "File to count lines in")
	linesCmd.MarkFlagRequired("file")

	rootCmd.AddCommand(countCmd, linesCmd)

	if err := rootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func countWord(cmd *cobra.Command, args []string) {
	filename, _ := cmd.Flags().GetString("file")
	word, _ := cmd.Flags().GetString("word")

	file, err := os.Open(filename)
	if err != nil {
		fmt.Printf("ファイルを開けませんでした: %v\n", err)
		os.Exit(1)
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	count := 0
	for scanner.Scan() {
		count += strings.Count(strings.ToLower(scanner.Text()), strings.ToLower(word))
	}

	if err := scanner.Err(); err != nil {
		fmt.Printf("ファイル読み込み中にエラーが発生しました: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("ファイル '%s' 内の '%s' の出現回数: %d\n", filename, word, count)
}

func countLines(cmd *cobra.Command, args []string) {
	filename, _ := cmd.Flags().GetString("file")

	file, err := os.Open(filename)
	if err != nil {
		fmt.Printf("ファイルを開けませんでした: %v\n", err)
		os.Exit(1)
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	lineCount := 0
	for scanner.Scan() {
		lineCount++
	}

	if err := scanner.Err(); err != nil {
		fmt.Printf("ファイル読み込み中にエラーが発生しました: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("ファイル '%s' の行数: %d\n", filename, lineCount)
}

動かしてみる

go build -o toolkit main.go

./toolkit count -f test.txt -w word test

ファイル 'test.txt' 内の 'word' の出現回数: 6

./toolkit lines -f test.txt

ファイル 'test.txt' の行数: 4

参考

• https://github.com/spf13/cobra