Update bash completion script (#1400)

* Update bash completion script * Filter config file extension
microsoft · Aug 5, 2019 · f179a9b · f179a9b
1 parent f23f8a0
commit f179a9b
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 40 deletions.
diff --git a/tools/bash-completion b/tools/bash-completion
@@ -1,30 +1,39 @@
 # list of commands/arguments
-__nnictl_cmds="create resume update stop trial experiment config webui log"
-__nnictl_create_cmds="--config --port"
-__nnictl_resume_cmds="--port"
+__nnictl_cmds="create resume update stop trial experiment platform import export webui config log package tensorboard top"
+__nnictl_create_cmds="--config --port --debug"
+__nnictl_resume_cmds="--port --debug"
 __nnictl_update_cmds="searchspace concurrency duration trialnum"
 __nnictl_update_searchspace_cmds="--filename"
 __nnictl_update_concurrency_cmds="--value"
 __nnictl_update_duration_cmds="--value"
 __nnictl_update_trialnum_cmds="--value"
-__nnictl_trial_cmds="ls kill"
-__nnictl_trial_kill_cmds="--trialid"
+__nnictl_stop_cmds="--port all"
+__nnictl_trial_cmds="ls kill codegen"
+__nnictl_trial_kill_cmds="--trial_id"
+__nnictl_trial_codegen_cmds="--trial_id"
+__nnictl_experiment_cmds="show status list delete"
+__nnictl_experiment_list_cmds="--all"
+__nnictl_experiment_delete_cmds="--all"
+__nnictl_platform_cmds="clean"
+__nnictl_platform_clean_cmds="--config"
+__nnictl_import_cmds="--filename"
+__nnictl_export_cmds="--type --filename"
 __nnictl_webui_cmds="url"
-__nnictl_experiment_cmds="show list status"
-__nnictl_experiment_list_cmds="all"
 __nnictl_config_cmds="show"
 __nnictl_log_cmds="stdout stderr trial"
 __nnictl_log_stdout_cmds="--tail --head --path"
 __nnictl_log_stderr_cmds="--tail --head --path"
+__nnictl_log_trial_cmds="--trial_id"
+__nnictl_package_cmds="install show"
+__nnictl_package_install_cmds="--name"
+__nnictl_tensorboard_cmds="start stop"
+__nnictl_tensorboard_start_cmds="--trial_id --port"
+__nnictl_top_cmds="--time"
 
-# list of arguments that accept a file name
-__nnictl_file_args=" --config -c --filename -f "
-
-# list of arguments that accept an experiment ID
-__nnictl_experiment_args=" --experiment -e "
-
-# list of arguments that accept a trial ID
-__nnictl_trial_args=" --trialid -t "
+# list of commands that accept an experiment ID as second argument
+__nnictl_2st_expid_cmds=" resume stop import export "
+# list of commands that accept an experiment ID as third argument
+__nnictl_3rd_expid_cmds=" update trial experiment webui config log tensorboard "
 
 
 # remove already set arguments from candidates
@@ -33,15 +42,27 @@ __nnictl_remain_args()
     local ret=${!1}  # ret = $__nnictl_xxx_cmds
     # for arg in COMP_WORDS[:-1]:
     for arg in "${COMP_WORDS[@]::${#COMP_WORDS[@]}-1}"; do
-        if [[ $arg == --* ]]; then
-            local ret=${ret/$arg/}  # remove it from $ret
-        fi
+        local ret=${ret/$arg/}  # remove it from $ret
     done
     echo $ret
 }
 
+# complete files with specific extension
+__nnictl_complete_extension()
+{
+    COMPREPLY=($(compgen -f -X "!*.$1" -- ${COMP_WORDS[-1]}))
+    if [[ -z "${COMPREPLY[*]}" ]]; then
+        # if there is no matching file here, search in sub-directories
+        COMPREPLY=($(compgen -d -S "/" -- ${COMP_WORDS[-1]}))
+        compopt -o nospace
+    fi
+}
+
 _nnictl()
 {
+    local cur=${COMP_WORDS[-1]}
+    local last=${COMP_WORDS[-2]}
+
     if [[ ${#COMP_WORDS[@]} -eq 2 ]]; then
         # completing frst argument from __nnictl_cmds
         COMPREPLY=($(compgen -W "$__nnictl_cmds" -- "${COMP_WORDS[1]}"))
@@ -51,33 +72,70 @@ _nnictl()
         local args=__nnictl_${COMP_WORDS[1]}_cmds
         COMPREPLY=($(compgen -W "${!args}" -- "${COMP_WORDS[2]}"))
 
-    elif [[ ${COMP_WORDS[-2]} != -* ]]; then
+        # add experiment IDs to candidates if desired
+        if [[ " resume stop import export " =~ " ${COMP_WORDS[1]} " ]]; then
+            local experiments=$(ls ~/nni/experiments 2>/dev/null)
+            COMPREPLY+=($(compgen -W "$experiments" -- $cur))
+        fi
+
+    elif [[ $last != -* || $last == --debug ]]; then
         # last argument does not starts with "-", so this one is likely to be "--xxx"
-        if [[ ${COMP_WORDS[2]} == -* ]]; then
-            # second argument starts with "-", use __nnictl_${FirstArg}_cmds
+        local args=__nnictl_${COMP_WORDS[1]}_${COMP_WORDS[2]}_cmds
+        if [[ $args =~ "-" || -z ${!args} ]]; then
+            # the second argument starts with "-", use __nnictl_${FirstArg}_cmds
             local args=__nnictl_${COMP_WORDS[1]}_cmds
-        else
-            # second argument is a word, use __nnictl_${FirstArg}_{SecondArg}_cmds
-            local args=__nnictl_${COMP_WORDS[1]}_${COMP_WORDS[2]}_cmds
         fi
         # remove already set arguments from candidates
         local remain_args=$(__nnictl_remain_args ${args})
-        COMPREPLY=($(compgen -W "$remain_args" -- "${COMP_WORDS[-1]}"))
+        COMPREPLY=($(compgen -W "$remain_args" -- $cur))
+
+        # if this is 3rd arguments, try adding experiment IDs to candidates
+        if [[ ${#COMP_WORDS[@]} -eq 4 ]]; then
+            if [[ $__nnictl_3rd_expid_cmds =~ " ${COMP_WORDS[1]} " && ${COMP_WORDS[2]} != "list" ]]; then
+                local experiments=$(ls ~/nni/experiments 2>/dev/null)
+                COMPREPLY+=($(compgen -W "$experiments" -- $cur))
+            fi
+        fi
+
+    elif [[ ${COMP_WORDS[1]} == "export" ]]; then
+        # "export" command is somewhat unique
+        if [[ " --type -t " =~ " $last " ]]; then
+            COMPREPLY=($(compgen -W "json csv" -- $cur))
+
+        elif [[ " --filename -f " =~ " $last " ]]; then
+            # try to detect whether complete CSV file or JSON file
+            [[ "$COMP_LINE" =~ "csv" ]] && local export_csv=1
+            [[ "$COMP_LINE" =~ "json" ]] && local export_json=1
+            if [[ -n $export_csv && -z $export_json ]]; then
+                local ext="csv"         # CSV only
+            elif [[ -n $export_json && -z $export_csv ]]; then
+                local ext="json"        # JSON only
+            else
+                local ext="@(csv|json)" # both
+            fi
+            __nnictl_complete_extension "$ext"
+        fi
+
+    elif [[ " --trial_id -t " =~ " $last " ]]; then
+        # complete trial ID
+        if [[ -e ${HOME}/nni/experiments/${COMP_WORDS[2]} ]]; then
+            local trials=$(ls -d ~/nni/experiments/${COMP_WORDS[2]}/trials/* 2>/dev/null | grep -o '[^/]*$')
+        elif [[ -e "${HOME}/nni/experiments/${COMP_WORDS[3]}" ]]; then
+            local trials=$(ls -d ~/nni/experiments/${COMP_WORDS[3]}/trials/* 2>/dev/null | grep -o '[^/]*$')
+        else
+            local trials=$(ls -d ~/nni/experiments/*/trials/* 2>/dev/null | grep -o '[^/]*$')
+        fi
+        COMPREPLY=($(compgen -W "$trials" -- $cur))
 
-    elif [[ $__nnictl_file_args =~ " ${COMP_WORDS[-2]} " ]]; then
-        # complete file names
-        COMPREPLY=($(compgen -f "${COMP_WORDS[-1]}"))
+    elif [[ " --config -c " =~ " $last " ]]; then
+        # complete YAML file
+        __nnictl_complete_extension "@(yml|yaml)"
 
-    elif [[ $__nnictl_experiment_args =~ " ${COMP_WORDS[-2]} " ]]; then
-        # complete experiment IDs
-        local experiments=$(ls ~/nni/experiments 2>/dev/null)
-        COMPREPLY=($(compgen -W "$experiments" -- "${COMP_WORDS[-1]}"))
+    elif [[ " --filename -f " =~ " $last " ]]; then
+        # complete JSON file
+        __nnictl_complete_extension "json"
 
-    elif [[ $__nnictl_trial_args =~ " ${COMP_WORDS[-2]} " ]]; then
-        # complete trial IDs
-        local trials=$(ls -d ~/nni/experiments/*/trials/* 2>/dev/null | grep -o '[^/]*$')
-        COMPREPLY=($(compgen -W "$trials" -- "${COMP_WORDS[-1]}"))
     fi
 }
 
-complete -o nospace -F _nnictl nnictl
+complete -o nosort -F _nnictl nnictl
diff --git a/tools/nni_cmd/nnictl.py b/tools/nni_cmd/nnictl.py
@@ -202,9 +202,9 @@ def parse_args():
     parser_tensorboard_start.add_argument('--trial_id', '-T', dest='trial_id', help='the id of trial')
     parser_tensorboard_start.add_argument('--port', dest='port', default=6006, help='the port to start tensorboard')
     parser_tensorboard_start.set_defaults(func=start_tensorboard)
-    parser_tensorboard_start = parser_tensorboard_subparsers.add_parser('stop', help='stop tensorboard')
-    parser_tensorboard_start.add_argument('id', nargs='?', help='the id of experiment')
-    parser_tensorboard_start.set_defaults(func=stop_tensorboard)
+    parser_tensorboard_stop = parser_tensorboard_subparsers.add_parser('stop', help='stop tensorboard')
+    parser_tensorboard_stop.add_argument('id', nargs='?', help='the id of experiment')
+    parser_tensorboard_stop.set_defaults(func=stop_tensorboard)
 
     #parse top command
     parser_top = subparsers.add_parser('top', help='monitor the experiment')