# list of commands/arguments

__nnictl_cmds="create resume view 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_view_cmds="--port"

__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_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_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 commands that accept an experiment ID as second argument

__nnictl_2nd_expid_cmds=" resume view 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

__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

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]}"))

elif [[ ${#COMP_WORDS[@]} -eq 3 ]]; then

# completing second argument from __nnictl_${FirstArg}_cmds

local args=__nnictl_${COMP_WORDS[1]}_cmds

COMPREPLY=($(compgen -W "${!args}" -- "${COMP_WORDS[2]}"))

# add experiment IDs to candidates if desired

if [[ $__nnictl_2nd_expid_cmds =~ " ${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"

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

fi

# remove already set arguments from candidates

local remain_args=$(__nnictl_remain_args ${args})

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 [[ " --config -c " =~ " $last " ]]; then

# complete YAML file

__nnictl_complete_extension "@(yml|yaml)"

elif [[ " --filename -f " =~ " $last " ]]; then

# complete JSON file

__nnictl_complete_extension "json"

fi

}

if [[ ${BASH_VERSINFO[0]} -le 4 && ${BASH_VERSINFO[1]} -le 4 ]]; then

complete -F _nnictl nnictl

else

complete -o nosort -F _nnictl nnictl

fi