Refactor flag parsing; start work on zsh completion

6 files changed, 46 insertions, 20 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 {