Add Atest tab completion script.

Now Atest can autocomplete module names, dir and filenames with
tab completion, and zsh users can benefit from it.

Bug: 110629292
Test: In both bash and zsh environments:
    source build/envsetup.sh && lunch aosp_arm-eng
    atest <tab><tab>  # has candidates
    adb <tab><tab>    # has candidates
    ENVSETUP_NO_COMPLETION=atest:adb . build/envsetup.sh && lunch aosp_arm-eng
    atest <tab><tab>  # no candidates
    adb <tab><tab>    # no candidates

Change-Id: Ib1c9e02feeb8aaf75c0b97821ae26e13ba8df350

1 files changed, 44 insertions, 26 deletions

diff --git a/envsetup.sh b/envsetup.sh
index 12168e122..58941440b 100644
--- a/envsetup.sh
+++ b/envsetup.sh
@@ -48,12 +48,12 @@ function build_build_var_cache()
 {
     local T=$(gettop)
     # Grep out the variable names from the script.
-    cached_vars=`cat $T/build/envsetup.sh | tr '()' '  ' | awk '{for(i=1;i<=NF;i++) if($i~/get_build_var/) print $(i+1)}' | sort -u | tr '\n' ' '`
-    cached_abs_vars=`cat $T/build/envsetup.sh | tr '()' '  ' | awk '{for(i=1;i<=NF;i++) if($i~/get_abs_build_var/) print $(i+1)}' | sort -u | tr '\n' ' '`
+    cached_vars=(`cat $T/build/envsetup.sh | tr '()' '  ' | awk '{for(i=1;i<=NF;i++) if($i~/get_build_var/) print $(i+1)}' | sort -u | tr '\n' ' '`)
+    cached_abs_vars=(`cat $T/build/envsetup.sh | tr '()' '  ' | awk '{for(i=1;i<=NF;i++) if($i~/get_abs_build_var/) print $(i+1)}' | sort -u | tr '\n' ' '`)
     # Call the build system to dump the "<val>=<value>" pairs as a shell script.
     build_dicts_script=`\builtin cd $T; build/soong/soong_ui.bash --dumpvars-mode \
-                        --vars="$cached_vars" \
-                        --abs-vars="$cached_abs_vars" \
+                        --vars="${cached_vars[*]}" \
+                        --abs-vars="${cached_abs_vars[*]}" \
                         --var-prefix=var_cache_ \
                         --abs-var-prefix=abs_var_cache_`
     local ret=$?
@@ -317,11 +317,11 @@ function set_sequence_number()
 
 # Takes a command name, and check if it's in ENVSETUP_NO_COMPLETION or not.
 function should_add_completion() {
-    local cmd="$1"
+    local cmd="$(basename $1| sed 's/_completion//' |sed 's/\.\(.*\)*sh$//')"
     case :"$ENVSETUP_NO_COMPLETION": in
-    *:"$cmd":*)
-        return 1
-        ;;
+        *:"$cmd":*)
+            return 1
+            ;;
     esac
     return 0
 }
@@ -330,22 +330,27 @@ function addcompletions()
 {
     local T dir f
 
-    # Keep us from trying to run in something that isn't bash.
-    if [ -z "${BASH_VERSION}" ]; then
+    # Keep us from trying to run in something that's neither bash nor zsh.
+    if [ -z "$BASH_VERSION" -a -z "$ZSH_VERSION" ]; then
         return
     fi
 
     # Keep us from trying to run in bash that's too old.
-    if [ ${BASH_VERSINFO[0]} -lt 3 ]; then
+    if [ -n "$BASH_VERSION" -a ${BASH_VERSINFO[0]} -lt 3 ]; then
         return
     fi
 
+    local completion_files=(
+      system/core/adb/adb.bash
+      system/core/fastboot/fastboot.bash
+      tools/tradefederation/core/atest/atest_completion.sh
+    )
     # Completion can be disabled selectively to allow users to use non-standard completion.
     # e.g.
     # ENVSETUP_NO_COMPLETION=adb # -> disable adb completion
     # ENVSETUP_NO_COMPLETION=adb:bit # -> disable adb and bit completion
-    for f in system/core/adb/adb.bash system/core/fastboot/fastboot.bash; do
-        if [ -f "$f" ] && should_add_completion $(basename "$f" .bash) ; then
+    for f in ${completion_files[*]}; do
+        if [ -f "$f" ] && should_add_completion "$f"; then
             . $f
         fi
     done
@@ -353,6 +358,7 @@ function addcompletions()
     if should_add_completion bit ; then
         complete -C "bit --tab" bit
     fi
+    complete -F _lunch lunch
 }
 
 function choosetype()
@@ -646,7 +652,6 @@ function _lunch()
     COMPREPLY=( $(compgen -W "${COMMON_LUNCH_CHOICES_CACHE}" -- ${cur}) )
     return 0
 }
-complete -F _lunch lunch
 
 # Configures the build to build unbundled apps.
 # Run tapas with one or more app names (from LOCAL_PACKAGE_NAME)
@@ -1562,24 +1567,37 @@ function atest()
     "$(gettop)"/tools/tradefederation/core/atest/atest.py "$@"
 }
 
-if [ "x$SHELL" != "x/bin/bash" ]; then
-    case `ps -o command -p $$` in
+# Zsh needs bashcompinit called to support bash-style completion.
+function add_zsh_completion() {
+    autoload -U compinit && compinit
+    autoload -U bashcompinit && bashcompinit
+}
+
+function validate_current_shell() {
+    local current_sh="$(ps -o command -p $$)"
+    case "$current_sh" in
         *bash*)
+            function check_type() { type -t "$1"; }
             ;;
+        *zsh*)
+            function check_type() { type "$1"; }
+            add_zsh_completion ;;
         *)
-            echo "WARNING: Only bash is supported, use of other shell would lead to erroneous results"
+            echo -e "WARNING: Only bash and zsh are supported.\nUse of other shell would lead to erroneous results."
             ;;
     esac
-fi
+}
 
 # Execute the contents of any vendorsetup.sh files we can find.
-for f in `test -d device && find -L device -maxdepth 4 -name 'vendorsetup.sh' 2> /dev/null | sort` \
-         `test -d vendor && find -L vendor -maxdepth 4 -name 'vendorsetup.sh' 2> /dev/null | sort` \
-         `test -d product && find -L product -maxdepth 4 -name 'vendorsetup.sh' 2> /dev/null | sort`
-do
-    echo "including $f"
-    . $f
-done
-unset f
+function source_vendorsetup() {
+    for dir in device vendor product; do
+        for f in $(test -d $dir && \
+            find -L $dir -maxdepth 4 -name 'vendorsetup.sh' 2>/dev/null | sort); do
+            echo "including $f"; . $f
+        done
+    done
+}
 
+validate_current_shell
+source_vendorsetup
 addcompletions