completion: decouple completion related stuff from command

It also now uses the main parser itself that invokes the command rather than
creating a new parser object.

dvc/cli/__init__.py

@@ -32,6 +32,7 @@ def parse_args(argv=None):
 
     parser = get_main_parser()
     args = parser.parse_args(argv)
+    args.parser = parser
     return args
 
 

dvc/cli/completion.py

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

dvc/commands/add.py

@@ -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__)
 

dvc/commands/cache.py

@@ -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
 

dvc/commands/check_ignore.py

@@ -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
 
 

dvc/commands/checkout.py

@@ -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
 

dvc/commands/commit.py

@@ -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__)
 

dvc/commands/completion.py

@@ -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

dvc/commands/daemon.py

@@ -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):

dvc/commands/data_sync.py

@@ -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__)
 

dvc/commands/diff.py

@@ -2,9 +2,9 @@ import argparse
 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__)

dvc/commands/experiments/apply.py

@@ -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__)
 

dvc/commands/experiments/diff.py

@@ -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

dvc/commands/experiments/push.py

@@ -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
 

dvc/commands/experiments/run.py

@@ -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

dvc/commands/experiments/show.py

@@ -8,9 +8,9 @@ from typing import TYPE_CHECKING
 
 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

dvc/commands/freeze.py

@@ -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__)

dvc/commands/get.py

@@ -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__)
 

dvc/commands/get_url.py

@@ -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__)
 

dvc/commands/imp.py

@@ -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__)

dvc/commands/imp_url.py

@@ -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__)

dvc/commands/ls/__init__.py

@@ -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

dvc/commands/metrics.py

@@ -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

dvc/commands/move.py

@@ -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__)

dvc/commands/params.py

@@ -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
 

dvc/commands/plots.py

@@ -5,9 +5,9 @@ import os
 
 from funcy import first
 
+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 import format_link

dvc/commands/remove.py

@@ -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__)

dvc/commands/repro.py

@@ -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.commands.status import CmdDataStatus
 
 

dvc/commands/stage.py

@@ -3,9 +3,9 @@ import logging
 from itertools import chain, filterfalse
 from typing import TYPE_CHECKING, Dict, Iterable, List
 
+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.utils.cli_parse import parse_params
 from dvc.utils.humanize import truncate_text
 

dvc/commands/unprotect.py

@@ -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__)

dvc/commands/update.py

@@ -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__)