From 0e1af655650a785a016e17ccd63c77e035838438 Mon Sep 17 00:00:00 2001
From: Zhenhui Xie <xiezh0831@yahoo.co.jp>
Date: Sun, 18 Oct 2020 12:58:44 +0800
Subject: [PATCH] Allow completing variables and recipes after `--set` in zsh
 completion script (#697)

---
 completions/just.zsh | 51 ++++++++++++++++++++++++++++++++++++----
 src/subcommand.rs    | 55 ++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 96 insertions(+), 10 deletions(-)

diff --git a/completions/just.zsh b/completions/just.zsh
index 0dda625..7ade921 100644
--- a/completions/just.zsh
+++ b/completions/just.zsh
@@ -64,15 +64,33 @@ _just() {
             curcontext="${curcontext%:*}-${words[2]}:"
 
             local lastarg=${words[${#words}]}
+            local recipe
 
-            if [[ ${lastarg} = */* ]]; then
+            local cmds; cmds=(
+                ${(s: :)$(_call_program commands just --summary)}
+            )
+
+            # Find first recipe name
+            for ((i = 2; i < $#words; i++ )) do
+                if [[ ${cmds[(I)${words[i]}]} -gt 0 ]]; then
+                    recipe=${words[i]}
+                    break
+                fi
+            done
+
+            if [[ $lastarg = */* ]]; then
                 # Arguments contain slash would be recognised as a file
                 _arguments -s -S $common '*:: :_files'
-            else
+            elif [[ $lastarg = *=* ]]; then
+                # Arguments contain equal would be recognised as a variable
+                _message "value"
+            elif [[ $recipe ]]; then
                 # Show usage message
-                _message "`just --show ${words[2]}`"
+                _message "`just --show $recipe`"
                 # Or complete with other commands
                 #_arguments -s -S $common '*:: :_just_commands'
+            else
+                _arguments -s -S $common '*:: :_just_commands'
             fi
         ;;
     esac
@@ -82,20 +100,43 @@ _just() {
 
 (( $+functions[_just_commands] )) ||
 _just_commands() {
+    [[ $PREFIX = -* ]] && return 1
+    integer ret=1
+    local variables; variables=(
+        ${(s: :)$(_call_program commands just --variables)}
+    )
     local commands; commands=(
         ${${${(M)"${(f)$(_call_program commands just --list)}":#    *}/ ##/}/ ##/:Args: }
     )
 
-    _describe -t commands 'just commands' commands "$@"
+    if compset -P '*='; then
+        case "${${words[-1]%=*}#*=}" in
+            *) _message 'value' && ret=0 ;;
+        esac
+    else
+        _describe -t variables 'variables' variables -qS "=" && ret=0
+        _describe -t commands 'just commands' commands "$@"
+    fi
+
 }
 
 (( $+functions[_just_variables] )) ||
 _just_variables() {
+    [[ $PREFIX = -* ]] && return 1
+    integer ret=1
     local variables; variables=(
         ${(s: :)$(_call_program commands just --variables)}
     )
 
-    _describe -t variables 'variables' variables
+    if compset -P '*='; then
+        case "${${words[-1]%=*}#*=}" in
+            *) _message 'value' && ret=0 ;;
+        esac
+    else
+        _describe -t variables 'variables' variables && ret=0
+    fi
+
+    return ret
 }
 
 _just "$@"
diff --git a/src/subcommand.rs b/src/subcommand.rs
index bcf8188..287e1de 100644
--- a/src/subcommand.rs
+++ b/src/subcommand.rs
@@ -73,15 +73,33 @@ const ZSH_COMPLETION_REPLACEMENTS: &[(&str, &str)] = &[
             curcontext="${curcontext%:*}-${words[2]}:"
 
             local lastarg=${words[${#words}]}
+            local recipe
 
-            if [[ ${lastarg} = */* ]]; then
+            local cmds; cmds=(
+                ${(s: :)$(_call_program commands just --summary)}
+            )
+
+            # Find first recipe name
+            for ((i = 2; i < $#words; i++ )) do
+                if [[ ${cmds[(I)${words[i]}]} -gt 0 ]]; then
+                    recipe=${words[i]}
+                    break
+                fi
+            done
+
+            if [[ $lastarg = */* ]]; then
                 # Arguments contain slash would be recognised as a file
                 _arguments -s -S $common '*:: :_files'
-            else
+            elif [[ $lastarg = *=* ]]; then
+                # Arguments contain equal would be recognised as a variable
+                _message "value"
+            elif [[ $recipe ]]; then
                 # Show usage message
-                _message "`just --show ${words[2]}`"
+                _message "`just --show $recipe`"
                 # Or complete with other commands
                 #_arguments -s -S $common '*:: :_just_commands'
+            else
+                _arguments -s -S $common '*:: :_just_commands'
             fi
         ;;
     esac
@@ -93,20 +111,47 @@ const ZSH_COMPLETION_REPLACEMENTS: &[(&str, &str)] = &[
     "    local commands; commands=(
 \x20\x20\x20\x20\x20\x20\x20\x20
     )",
-    r#"    local commands; commands=(
+    r#"    [[ $PREFIX = -* ]] && return 1
+    integer ret=1
+    local variables; variables=(
+        ${(s: :)$(_call_program commands just --variables)}
+    )
+    local commands; commands=(
         ${${${(M)"${(f)$(_call_program commands just --list)}":#    *}/ ##/}/ ##/:Args: }
     )
+"#,
+  ),
+  (
+    r#"    _describe -t commands 'just commands' commands "$@""#,
+    r#"    if compset -P '*='; then
+        case "${${words[-1]%=*}#*=}" in
+            *) _message 'value' && ret=0 ;;
+        esac
+    else
+        _describe -t variables 'variables' variables -qS "=" && ret=0
+        _describe -t commands 'just commands' commands "$@"
+    fi
 "#,
   ),
   (
     r#"_just "$@""#,
     r#"(( $+functions[_just_variables] )) ||
 _just_variables() {
+    [[ $PREFIX = -* ]] && return 1
+    integer ret=1
     local variables; variables=(
         ${(s: :)$(_call_program commands just --variables)}
     )
 
-    _describe -t variables 'variables' variables
+    if compset -P '*='; then
+        case "${${words[-1]%=*}#*=}" in
+            *) _message 'value' && ret=0 ;;
+        esac
+    else
+        _describe -t variables 'variables' variables && ret=0
+    fi
+
+    return ret
 }
 
 _just "$@""#,