From a7d37b80657ff37c15cf9e68575e8e9c873f815d Mon Sep 17 00:00:00 2001
From: Paul Ouellette <oue.paul18@gmail.com>
Date: Fri, 26 Jul 2019 23:31:09 -0400
Subject: [PATCH] Zsh completions: complete main command options after
 subcommands

---
 spec/completion_spec.lua | 57 +++++++++++++++++++++++++++++++---------
 src/argparse.lua         | 14 +++++++---
 2 files changed, 55 insertions(+), 16 deletions(-)

diff --git a/spec/completion_spec.lua b/spec/completion_spec.lua
index 6812cd2..94b3927 100644
--- a/spec/completion_spec.lua
+++ b/spec/completion_spec.lua
@@ -92,63 +92,94 @@ _comptest() {
   local context state state_descr line
   typeset -A opt_args
 
+  local -a options=(
+    {-h,--help}"[Show this help message and exit]"
+    "--completion[Output a shell completion script for the specified shell]: :(bash zsh fish)"
+    "*"{-v,--verbose}"[Set the verbosity level]"
+    {-f,--files}"[A description with illegal \"' characters]:*: :_files"
+  )
   _arguments -s -S \
-    {-h,--help}"[Show this help message and exit]" \
-    "--completion[Output a shell completion script for the specified shell]: :(bash zsh fish)" \
-    "*"{-v,--verbose}"[Set the verbosity level]" \
-    {-f,--files}"[A description with illegal \"' characters]:*: :_files" \
+    $options \
     ": :_comptest_cmds" \
     "*:: :->args" \
     && return 0
 
   case $words[1] in
     help)
+      options=(
+        $options
+        {-h,--help}"[Show this help message and exit]"
+      )
       _arguments -s -S \
-        {-h,--help}"[Show this help message and exit]" \
+        $options \
         ": :(help completion install i admin)" \
         && return 0
       ;;
 
     completion)
+      options=(
+        $options
+        {-h,--help}"[Show this help message and exit]"
+      )
       _arguments -s -S \
-        {-h,--help}"[Show this help message and exit]" \
+        $options \
         ": :(bash zsh fish)" \
         && return 0
       ;;
 
     install|i)
+      options=(
+        $options
+        {-h,--help}"[Show this help message and exit]"
+        "--deps-mode: :(all one order none)"
+        "--no-doc[Install without documentation]"
+      )
       _arguments -s -S \
-        {-h,--help}"[Show this help message and exit]" \
-        "--deps-mode: :(all one order none)" \
-        "--no-doc[Install without documentation]" \
+        $options \
         && return 0
       ;;
 
     admin)
+      options=(
+        $options
+        {-h,--help}"[Show this help message and exit]"
+      )
       _arguments -s -S \
-        {-h,--help}"[Show this help message and exit]" \
+        $options \
         ": :_comptest_admin_cmds" \
         "*:: :->args" \
         && return 0
 
       case $words[1] in
         help)
+          options=(
+            $options
+            {-h,--help}"[Show this help message and exit]"
+          )
           _arguments -s -S \
-            {-h,--help}"[Show this help message and exit]" \
+            $options \
             ": :(help add remove)" \
             && return 0
           ;;
 
         add)
+          options=(
+            $options
+            {-h,--help}"[Show this help message and exit]"
+          )
           _arguments -s -S \
-            {-h,--help}"[Show this help message and exit]" \
+            $options \
             ": :_files" \
             && return 0
           ;;
 
         remove)
+          options=(
+            $options
+            {-h,--help}"[Show this help message and exit]"
+          )
           _arguments -s -S \
-            {-h,--help}"[Show this help message and exit]" \
+            $options \
             ": :_files" \
             && return 0
           ;;
diff --git a/src/argparse.lua b/src/argparse.lua
index d105d43..59c9be3 100644
--- a/src/argparse.lua
+++ b/src/argparse.lua
@@ -1306,7 +1306,12 @@ complete -F _%s -o bashdefault -o default %s
 end
 
 function Parser:_zsh_arguments(buf, cmd_name, indent)
-   table.insert(buf, (" "):rep(indent) .. "_arguments -s -S \\")
+   if self._parent then
+      table.insert(buf, (" "):rep(indent) .. "options=(")
+      table.insert(buf, (" "):rep(indent + 2) .. "$options")
+   else
+      table.insert(buf, (" "):rep(indent) .. "local -a options=(")
+   end
 
    for _, option in ipairs(self._options) do
       local line = {}
@@ -1335,10 +1340,13 @@ function Parser:_zsh_arguments(buf, cmd_name, indent)
          table.insert(line, ": :_files")
       end
       table.insert(line, '"')
-
-      table.insert(buf, (" "):rep(indent + 2) .. table.concat(line) .. " \\")
+      table.insert(buf, (" "):rep(indent + 2) .. table.concat(line))
    end
 
+   table.insert(buf, (" "):rep(indent) .. ")")
+   table.insert(buf, (" "):rep(indent) .. "_arguments -s -S \\")
+   table.insert(buf, (" "):rep(indent + 2) .. "$options \\")
+
    if self._is_help_command then
       table.insert(buf, (" "):rep(indent + 2) .. '": :(' .. self._parent:_get_commands() .. ')" \\')
    else