Skip to content

Commit

Permalink
completion: decouple completion related stuff from command
Browse files Browse the repository at this point in the history
It also now uses the main parser itself that invokes the command rather than
creating a new parser object.
  • Loading branch information
skshetry committed Apr 20, 2022
1 parent 5c1e7be commit cc8a0ce
Show file tree
Hide file tree
Showing 31 changed files with 105 additions and 103 deletions.
1 change: 1 addition & 0 deletions dvc/cli/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ def parse_args(argv=None):

parser = get_main_parser()
args = parser.parse_args(argv)
args.parser = parser
return args


Expand Down
73 changes: 73 additions & 0 deletions dvc/cli/completion.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
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
}
"""

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

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"}
2 changes: 1 addition & 1 deletion dvc/commands/add.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion

logger = logging.getLogger(__name__)

Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/cache.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import argparse

from dvc.cli import completion
from dvc.cli.utils import append_doc_link, fix_subparsers
from dvc.commands import completion
from dvc.commands.config import CmdConfig
from dvc.ui import ui

Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/check_ignore.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import argparse

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.ui import ui


Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/checkout.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import operator

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.exceptions import CheckoutError
from dvc.ui import ui

Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/commit.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion

logger = logging.getLogger(__name__)

Expand Down
78 changes: 3 additions & 75 deletions dvc/commands/completion.py
Original file line number Diff line number Diff line change
@@ -1,92 +1,20 @@
import argparse
import logging

import shtab

from dvc.cli.command import CmdBaseNoRepo
from dvc.cli.completion import PREAMBLE
from dvc.cli.utils import append_doc_link
from dvc.ui import ui

logger = logging.getLogger(__name__)
FILE = shtab.FILE
DIR = shtab.DIRECTORY

PREAMBLE = {
"bash": """
# $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
}
""",
"zsh": """
_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_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"}


class CmdCompletion(CmdBaseNoRepo):
def run(self):
from dvc.cli.parser import get_main_parser
import shtab

parser = get_main_parser()
shell = self.args.shell
parser = self.args.parser
script = shtab.complete(parser, shell=shell, preamble=PREAMBLE)
ui.write(script, force=True)
return 0
Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/daemon.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from dvc.cli import completion
from dvc.cli.command import CmdBaseNoRepo
from dvc.cli.utils import fix_subparsers
from dvc.commands import completion


class CmdDaemonBase(CmdBaseNoRepo):
Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/data_sync.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion

logger = logging.getLogger(__name__)

Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/diff.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@
import logging
import os

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.ui import ui

logger = logging.getLogger(__name__)
Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/experiments/apply.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion

logger = logging.getLogger(__name__)

Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/experiments/diff.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.commands.metrics import DEFAULT_PRECISION
from dvc.exceptions import DvcException
from dvc.ui import ui
Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/experiments/push.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.exceptions import InvalidArgumentError
from dvc.ui import ui

Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/experiments/run.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.commands.repro import CmdRepro
from dvc.commands.repro import add_arguments as add_repro_arguments
from dvc.exceptions import InvalidArgumentError
Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/experiments/show.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,9 @@

from funcy import lmap

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.commands.metrics import DEFAULT_PRECISION
from dvc.exceptions import DvcException, InvalidArgumentError
from dvc.ui import ui
Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/freeze.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.exceptions import DvcException

logger = logging.getLogger(__name__)
Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/get.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBaseNoRepo
from dvc.exceptions import DvcException

from ..cli.utils import append_doc_link
from . import completion

logger = logging.getLogger(__name__)

Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/get_url.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBaseNoRepo
from dvc.exceptions import DvcException

from ..cli.utils import append_doc_link
from . import completion

logger = logging.getLogger(__name__)

Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/imp.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.exceptions import DvcException

logger = logging.getLogger(__name__)
Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/imp_url.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.exceptions import DvcException

logger = logging.getLogger(__name__)
Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/ls/__init__.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBaseNoRepo
from dvc.cli.utils import append_doc_link
from dvc.commands import completion
from dvc.commands.ls.ls_colors import LsColors
from dvc.exceptions import DvcException
from dvc.ui import ui
Expand Down
2 changes: 1 addition & 1 deletion dvc/commands/metrics.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import argparse
import logging

from dvc.cli import completion
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link, fix_subparsers
from dvc.commands import completion
from dvc.exceptions import DvcException
from dvc.ui import ui
from dvc.utils.serialize import encode_exception
Expand Down
Loading

0 comments on commit cc8a0ce

Please sign in to comment.