Refactor flag parsing; start work on zsh completion

7 files changed, 79 insertions, 37 deletions

diff --git a/cmd/completions.go b/cmd/completions.go
index e169952..285f235 100644
--- a/cmd/completions.go
+++ b/cmd/completions.go
@@ -1,7 +1,13 @@
 package cmd
 
-import "fmt"
+import (
+	_ "embed"
+	"fmt"
+)
+
+//go:embed completions/atlas.zsh
+var zshCompletions string
 
 func ZshCompletions() {
-	fmt.Println("Not implemented yet!")
+	fmt.Print(zshCompletions)
 }
diff --git a/cmd/completions/atlas.zsh b/cmd/completions/atlas.zsh
new file mode 100644
index 0000000..c9f45e0
--- /dev/null
+++ b/cmd/completions/atlas.zsh
@@ -0,0 +1,17 @@
+#compdef atlas
+
+_atlas() {
+	local -a commands
+	commands=(
+		'query:Query stuff'
+		'index:Index stuff'
+		'server:Launch server'
+		'shell:Launch interactive shell'
+		'completions:Generate shell completions'
+		'help:View help'
+	)
+
+	_describe 'command' commands
+}
+
+_atlas "$@"
diff --git a/cmd/help.go b/cmd/help.go
index a389cb8..e566448 100644
--- a/cmd/help.go
+++ b/cmd/help.go
@@ -18,6 +18,7 @@ var helpTopics = []string{
 	"query", "q",
 	"shell",
 	"server",
+	"completions",
 }
 
 func PrintHelp(w io.Writer) {
@@ -28,6 +29,7 @@ func PrintHelp(w io.Writer) {
 	fmt.Fprintln(w, "  query <subcommand>    - search against an index")
 	fmt.Fprintln(w, "  shell                 - start a debug shell")
 	fmt.Fprintln(w, "  server                - start an http query server (EXPERIMENTAL)")
+	fmt.Fprintln(w, "  completions           - shell completions (EXPERIMENTAL)")
 	fmt.Fprintln(w, "  help  <help-topic>    - print help info")
 }
 
@@ -44,10 +46,9 @@ func PrintFlagSet(w io.Writer, fs *flag.FlagSet) {
 }
 
 func Help(topic string, w io.Writer) {
-	fs := flag.NewFlagSet(topic, flag.ExitOnError)
 	switch topic {
 	case "index", "i":
-		SetupIndexFlags(nil, fs, &IndexFlags{})
+		fs := NewIndexFlagSet(&IndexFlags{})
 		fmt.Fprintf(w, "%s [global-flags] index [index-flags] <subcommand>\n\n", os.Args[0])
 		fmt.Fprintln(w, "Subcommands:")
 		fmt.Fprintln(w, "  build  - create a new index")
@@ -69,7 +70,7 @@ func Help(topic string, w io.Writer) {
 		fmt.Fprintf(w, "%s [global-flags] index tidy\n\n", os.Args[0])
 		fmt.Fprintln(w, "Remove unused authors or tags and optimize the database")
 	case "query", "q":
-		SetupQueryFlags(nil, fs, &QueryFlags{}, "")
+		fs := NewQueryFlagSet(&QueryFlags{}, "")
 		fmt.Fprintf(w, "%s [global-flags] query [query-flags] <query>...\n\n", os.Args[0])
 		fmt.Fprintln(w, "Execute a query against the connected database")
 		fmt.Fprintln(w, "Query Flags:")
@@ -156,7 +157,7 @@ Atlas recognizes many of the common date formats.
 		fmt.Fprintln(w, "\nShell Help:")
 		shell.PrintHelp(w)
 	case "server":
-		SetupServerFlags(nil, fs, &ServerFlags{})
+		fs := NewServerFlagSet(&ServerFlags{})
 		fmt.Fprintf(w, "%s [global-flags] server [server-flags]\n", os.Args[0])
 		fmt.Fprintln(w, "Run a server to execute queries over HTTP or a unix domain socket")
 		fmt.Fprintln(w, "HTTP Server:")
@@ -167,6 +168,11 @@ Atlas recognizes many of the common date formats.
 		fmt.Fprintln(w, "    sortOrder: desc, descending")
 		fmt.Fprintln(w, "Server Flags:")
 		PrintFlagSet(w, fs)
+	case "completions":
+		fmt.Fprintf(w, "%s completions <shell>\n", os.Args[0])
+		fmt.Fprintln(w, "Generate shell completions")
+		fmt.Fprintln(w, "Supported Shells:")
+		fmt.Fprintln(w, "  zsh")
 	case "help", "":
 		PrintHelp(w)
 		fmt.Fprintln(w, "\nHelp Topics:")
diff --git a/cmd/index.go b/cmd/index.go
index ba9bb8f..7fd63ab 100644
--- a/cmd/index.go
+++ b/cmd/index.go
@@ -19,7 +19,8 @@ type IndexFlags struct {
 	index.ParseOpts
 }
 
-func SetupIndexFlags(args []string, fs *flag.FlagSet, flags *IndexFlags) {
+func NewIndexFlagSet(flags *IndexFlags) *flag.FlagSet {
+	fs := flag.NewFlagSet("index", flag.ExitOnError)
 	flags.ParseLinks = true
 	flags.ParseMeta = true
 	flags.ParseHeadings = true
@@ -65,14 +66,7 @@ func SetupIndexFlags(args []string, fs *flag.FlagSet, flags *IndexFlags) {
 		PrintGlobalFlags(f)
 	}
 
-	fs.Parse(args)
-
-	remainingArgs := fs.Args()
-	if len(remainingArgs) == 0 {
-		flags.Subcommand = "build"
-	} else if len(remainingArgs) == 1 {
-		flags.Subcommand = remainingArgs[0]
-	}
+	return fs
 }
 
 func RunIndex(gFlags GlobalFlags, iFlags IndexFlags, db *data.Query) byte {
diff --git a/cmd/query.go b/cmd/query.go
index eeb67ac..7178c0d 100644
--- a/cmd/query.go
+++ b/cmd/query.go
@@ -22,7 +22,10 @@ type QueryFlags struct {
 	SortDesc          bool
 }
 
-func SetupQueryFlags(args []string, fs *flag.FlagSet, flags *QueryFlags, dateFormat string) {
+func NewQueryFlagSet(flags *QueryFlags, dateFormat string) *flag.FlagSet {
+
+	fs := flag.NewFlagSet("query", flag.ExitOnError)
+
 	// NOTE: providing `-outFormat` before `-outCustomFormat` might ignore user specified format
 	fs.Func("outFormat", "output `format` for queries (default, json, yaml, pathonly, custom)",
 		func(arg string) error {
@@ -63,8 +66,7 @@ func SetupQueryFlags(args []string, fs *flag.FlagSet, flags *QueryFlags, dateFor
 		PrintFlagSet(w, fs)
 		PrintGlobalFlags(w)
 	}
-
-	fs.Parse(args)
+	return fs
 }
 
 func RunQuery(gFlags GlobalFlags, qFlags QueryFlags, db *data.Query, searchQuery string) byte {
diff --git a/cmd/server.go b/cmd/server.go
index 9de8dbc..fe80346 100644
--- a/cmd/server.go
+++ b/cmd/server.go
@@ -21,11 +21,12 @@ type ServerFlags struct {
 	Port    int
 }
 
-func SetupServerFlags(args []string, fs *flag.FlagSet, flags *ServerFlags) {
+func NewServerFlagSet(flags *ServerFlags) *flag.FlagSet {
+	fs := flag.NewFlagSet("server", flag.ExitOnError)
 	fs.StringVar(&flags.Address, "address", "127.0.0.1", "the address to listen on, prefix with 'unix:' to create a unixsocket")
 	fs.IntVar(&flags.Port, "port", 8080, "the port to bind to")
 
-	fs.Parse(args)
+	return fs
 }
 
 func RunServer(gFlags GlobalFlags, sFlags ServerFlags, db *data.Query) byte {
diff --git a/main.go b/main.go
index 4329a61..eb012a8 100644
--- a/main.go
+++ b/main.go
@@ -26,26 +26,29 @@ func addGlobalFlagUsage(fs *flag.FlagSet) func() {
 }
 
 func main() {
+	exitCode := 0
+	defer func() { os.Exit(exitCode) }()
+
+	// global flags
 	globalFlags := cmd.GlobalFlags{}
 	cmd.SetupGlobalFlags(flag.CommandLine, &globalFlags)
+	flag.Parse()
+	args := flag.Args()
 
-	indexFs := flag.NewFlagSet("index", flag.ExitOnError)
-	queryFs := flag.NewFlagSet("query", flag.ExitOnError)
+	queryFlags := cmd.QueryFlags{Outputer: query.DefaultOutput{}}
+	indexFlags := cmd.IndexFlags{}
+	serverFlags := cmd.ServerFlags{Port: 8080}
+
+	indexFs := cmd.NewIndexFlagSet(&indexFlags)
+	queryFs := cmd.NewQueryFlagSet(&queryFlags, globalFlags.DateFormat)
 	shellFs := flag.NewFlagSet("debug", flag.ExitOnError)
-	serverFs := flag.NewFlagSet("server", flag.ExitOnError)
+	serverFs := cmd.NewServerFlagSet(&serverFlags)
 	completionsFs := flag.NewFlagSet("completions", flag.ContinueOnError)
 
 	// set default usage for flagsets without subcommands
 	shellFs.Usage = addGlobalFlagUsage(shellFs)
 	serverFs.Usage = addGlobalFlagUsage(serverFs)
 
-	flag.Parse()
-	args := flag.Args()
-
-	queryFlags := cmd.QueryFlags{Outputer: query.DefaultOutput{}}
-	indexFlags := cmd.IndexFlags{}
-	serverFlags := cmd.ServerFlags{Port: 8080}
-
 	if len(args) < 1 {
 		fmt.Fprintln(os.Stderr, "No Command provided")
 		cmd.PrintHelp(os.Stderr)
@@ -56,11 +59,17 @@ func main() {
 
 	switch command {
 	case "query", "q":
-		cmd.SetupQueryFlags(args[1:], queryFs, &queryFlags, globalFlags.DateFormat)
+		queryFs.Parse(args[1:])
 	case "index", "i":
-		cmd.SetupIndexFlags(args[1:], indexFs, &indexFlags)
+		indexFs.Parse(args[1:])
+		remainingArgs := indexFs.Args()
+		if len(remainingArgs) == 0 {
+			indexFlags.Subcommand = "build"
+		} else if len(remainingArgs) == 1 {
+			indexFlags.Subcommand = remainingArgs[0]
+		}
 	case "server":
-		cmd.SetupServerFlags(args[1:], serverFs, &serverFlags)
+		serverFs.Parse(args[1:])
 	case "completions":
 		completionsFs.Parse(args[1:])
 	case "help":
@@ -126,10 +135,11 @@ func main() {
 	logger := slog.New(logHandler)
 	slog.SetDefault(logger)
 
+	// TODO: remove later querier instances
 	querier := data.NewQuery(globalFlags.DBPath, VERSION)
+	defer querier.Close()
 
 	// command specific
-	var exitCode int
 	switch command {
 	case "query", "q":
 		searchQuery := strings.Join(queryFs.Args(), " ")
@@ -140,6 +150,15 @@ func main() {
 		exitCode = int(cmd.RunServer(globalFlags, serverFlags, querier))
 	case "completions":
 		lang := completionsFs.Arg(0)
+		fmt.Fprintln(os.Stderr, "General Flags")
+		flag.VisitAll(func(f *flag.Flag) {
+			fmt.Fprintf(os.Stderr, "%s - %s\n", f.Name, f.Usage)
+		})
+		fmt.Fprintln(os.Stderr, "Index Flags")
+		indexFs.VisitAll(func(f *flag.Flag) {
+			fmt.Fprintf(os.Stderr, "%s - %s\n", f.Name, f.Usage)
+		})
+
 		switch lang {
 		case "zsh":
 			cmd.ZshCompletions()
@@ -152,7 +171,4 @@ func main() {
 	case "shell":
 		exitCode = int(cmd.RunShell(globalFlags, querier, VERSION))
 	}
-
-	querier.Close()
-	os.Exit(exitCode)
 }