mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-10 15:28:34 +01:00
110 lines
2.8 KiB
Python
110 lines
2.8 KiB
Python
import shtab
|
|
|
|
BASH_PREAMBLE = """
|
|
# $1=COMP_WORDS[1]
|
|
_dvc_compgen_DVCFiles() {
|
|
compgen -d -S '/' -- $1 # recurse into subdirs
|
|
compgen -f -X '!*?.dvc' -- $1
|
|
compgen -f -X '!*Dvcfile' -- $1
|
|
compgen -f -X '!*dvc.yaml' -- $1
|
|
}
|
|
|
|
_dvc_compgen_stages() {
|
|
local _dvc_stages=($(dvc stage list -q --names-only))
|
|
compgen -W "${_dvc_stages[*]}" -- $1
|
|
}
|
|
_dvc_compgen_stages_and_files() {
|
|
_dvc_compgen_DVCFiles $1
|
|
_dvc_compgen_stages $1
|
|
}
|
|
|
|
_dvc_compgen_exps() {
|
|
local _dvc_exps=($(dvc exp list -q --all-commits --names-only))
|
|
compgen -W "${_dvc_exps[*]}" -- $1
|
|
}
|
|
|
|
_dvc_compgen_remotes() {
|
|
local _dvc_remotes=($(dvc remote list | cut -d' ' -f1))
|
|
compgen -W "${_dvc_remotes[*]}" -- $1
|
|
}
|
|
|
|
_dvc_compgen_config_vars() {
|
|
compgen -W "${_dvc_config_vars[*]}" -- $1
|
|
}
|
|
"""
|
|
|
|
ZSH_PREAMBLE = """
|
|
_dvc_compadd_DVCFiles() {
|
|
_files -g '(*?.dvc|Dvcfile|dvc.yaml)'
|
|
}
|
|
_dvc_compadd_stages() {
|
|
# this will also show up the description of the stages
|
|
_describe 'stages' "($(_dvc_stages_output))"
|
|
}
|
|
|
|
_dvc_stages_output() {
|
|
dvc stage list -q | awk '{
|
|
# escape possible `:` on the stage name
|
|
sub(/:/, "\\\\\\\\:", $1);
|
|
# read all of the columns except the first
|
|
# reading `out` from $2, so as not to have a leading whitespace
|
|
out=$2; for(i=3;i<=NF;i++){out=out" "$i};
|
|
# print key, ":" and then single-quote the description
|
|
# colon is a delimiter used by `_describe` to separate field/description
|
|
print $1":""\\047"out"\\047"
|
|
# single quote -> \\047
|
|
}'
|
|
}
|
|
|
|
_dvc_compadd_stages_and_files() {
|
|
_dvc_compadd_DVCFiles
|
|
_dvc_compadd_stages
|
|
}
|
|
|
|
_dvc_compadd_exps() {
|
|
_describe 'experiments' "($(dvc exp list -q --all-commits --names-only))"
|
|
}
|
|
|
|
_dvc_compadd_remotes() {
|
|
_describe 'remotes' "($(dvc remote list | cut -d' ' -f1))"
|
|
}
|
|
|
|
_dvc_compadd_config_vars() {
|
|
_describe 'config_vars' _dvc_config_vars
|
|
}
|
|
"""
|
|
|
|
PREAMBLE = {
|
|
"bash": BASH_PREAMBLE,
|
|
"zsh": ZSH_PREAMBLE,
|
|
}
|
|
|
|
FILE = shtab.FILE
|
|
DIR = shtab.DIRECTORY
|
|
DVC_FILE = {"bash": "_dvc_compgen_DVCFiles", "zsh": "_dvc_compadd_DVCFiles"}
|
|
STAGE = {"bash": "_dvc_compgen_stages", "zsh": "_dvc_compadd_stages"}
|
|
DVCFILES_AND_STAGE = {
|
|
"bash": "_dvc_compgen_stages_and_files",
|
|
"zsh": "_dvc_compadd_stages_and_files",
|
|
}
|
|
EXPERIMENT = {"bash": "_dvc_compgen_exps", "zsh": "_dvc_compadd_exps"}
|
|
REMOTE = {"bash": "_dvc_compgen_remotes", "zsh": "_dvc_compadd_remotes"}
|
|
CONFIG_VARS = {"bash": "_dvc_compgen_config_vars", "zsh": "_dvc_compadd_config_vars"}
|
|
|
|
|
|
def get_preamble() -> dict[str, str]:
|
|
from dvc.config_schema import config_vars_for_completion
|
|
|
|
ret: dict[str, str] = {}
|
|
config_vars = list(config_vars_for_completion())
|
|
|
|
nl = "\n\t".expandtabs(4)
|
|
config_vars_arr = f"""
|
|
_dvc_config_vars=(
|
|
{nl.join(config_vars)}
|
|
)
|
|
"""
|
|
for shell, preamble in PREAMBLE.items():
|
|
ret[shell] = config_vars_arr + preamble
|
|
return ret
|