diff --git a/dvc/command/completion.py b/dvc/command/completion.py index 644a411ed..c6fae384d 100644 --- a/dvc/command/completion.py +++ b/dvc/command/completion.py @@ -8,10 +8,7 @@ from dvc.command.base import CmdBaseNoRepo, append_doc_link logger = logging.getLogger(__name__) FILE = shtab.FILE DIR = shtab.DIRECTORY -DVC_FILE = { - "bash": "_dvc_compgen_DVCFiles", - "zsh": "_files -g '(*?.dvc|Dvcfile|dvc.yaml)'", -} + PREAMBLE = { "bash": """ # $1=COMP_WORDS[1] @@ -21,8 +18,59 @@ _dvc_compgen_DVCFiles() { compgen -f -X '!*Dvcfile' -- $1 compgen -f -X '!*dvc.yaml' -- $1 } -""", - "zsh": "", + +_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 +} + """, + "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_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", } diff --git a/dvc/command/repro.py b/dvc/command/repro.py index e395546c8..335130beb 100644 --- a/dvc/command/repro.py +++ b/dvc/command/repro.py @@ -50,7 +50,7 @@ def add_arguments(repro_parser): "targets", nargs="*", help="Stages to reproduce. 'dvc.yaml' by default.", - ).complete = completion.DVC_FILE + ).complete = completion.DVCFILES_AND_STAGE repro_parser.add_argument( "-f", "--force", diff --git a/dvc/command/stage.py b/dvc/command/stage.py index c2463804a..7ca1355d7 100644 --- a/dvc/command/stage.py +++ b/dvc/command/stage.py @@ -1,11 +1,109 @@ import argparse import logging +from itertools import chain, filterfalse +from typing import TYPE_CHECKING, Dict, Iterable, List from dvc.command import completion from dvc.command.base import CmdBase, append_doc_link, fix_subparsers +if TYPE_CHECKING: + from dvc.output.base import BaseOutput + from dvc.stage import Stage + logger = logging.getLogger(__name__) +MAX_TEXT_LENGTH = 80 +ELLIPSIS = "…" + + +def generate_description(stage: "Stage") -> str: + def part_desc(outs: Iterable["BaseOutput"]) -> str: + return ", ".join(out.def_path for out in outs) + + if not stage.deps and not stage.outs: + return "No outputs or dependencies." + + if not stage.outs and stage.deps: + return "Stage depends on " + part_desc(stage.deps) + + def is_plot_or_metric(out: "BaseOutput"): + return bool(out.plot) or bool(out.metric) + + desc: List[str] = [] + + outs = list(filterfalse(is_plot_or_metric, stage.outs)) + if outs: + desc.append("Produces " + part_desc(outs)) + + plots_and_metrics = list(filter(is_plot_or_metric, stage.outs)) + if plots_and_metrics: + desc.append("Generates " + part_desc(plots_and_metrics)) + + return "; ".join(desc) + + +def prepare_description( + stage: "Stage", max_length: int = MAX_TEXT_LENGTH +) -> str: + desc = stage.short_description() or generate_description(stage) + if len(desc) > max_length: + return desc[: max_length - 1] + ELLIPSIS + return desc + + +def prepare_stages_data( + stages: Iterable["Stage"], + description: bool = True, + max_length: int = MAX_TEXT_LENGTH, +) -> Dict[str, str]: + return { + stage.addressing: prepare_description(stage, max_length=max_length) + if description + else "" + for stage in stages + } + + +class CmdStageList(CmdBase): + def _get_stages(self) -> Iterable["Stage"]: + if self.args.all: + stages: List["Stage"] = self.repo.stages # type: ignore + logger.trace( # type: ignore[attr-defined] + "%d no. of stages found", len(stages) + ) + return stages + + # removing duplicates while maintaining order + collected = chain.from_iterable( + self.repo.stage.collect( + target=target, + recursive=self.args.recursive, + accept_group=True, + ) + for target in self.args.targets + ) + return dict.fromkeys(collected).keys() + + def run(self): + from dvc.utils.diff import table + + def log_error(relpath: str, exc: Exception): + if self.args.fail: + raise exc + logger.debug("Stages from %s failed to load", relpath) + + # silence stage collection error by default + self.repo.stage_collection_error_handler = log_error + + stages = self._get_stages() + names_only = self.args.names_only + + data = prepare_stages_data(stages, description=not names_only) + if data: + print(table(header=(), rows=data.items())) + + return 0 + class CmdStageAdd(CmdBase): def run(self): @@ -194,3 +292,47 @@ def add_parser(subparsers, parent_parser): ) _add_common_args(stage_add_parser) stage_add_parser.set_defaults(func=CmdStageAdd) + + STAGE_LIST_HELP = "List stages." + stage_list_parser = stage_subparsers.add_parser( + "list", + parents=[parent_parser], + description=append_doc_link(STAGE_LIST_HELP, "stage/list"), + help=STAGE_LIST_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + stage_list_parser.add_argument( + "targets", + nargs="*", + default=["dvc.yaml"], + help=( + "Show stages from a dvc.yaml/.dvc file or a directory. " + "'dvc.yaml' by default" + ), + ) + stage_list_parser.add_argument( + "--all", + action="store_true", + default=False, + help="List all of the stages in the repo.", + ) + stage_list_parser.add_argument( + "--fail", + action="store_true", + default=False, + help="Fail immediately if there's an error.", + ) + stage_list_parser.add_argument( + "-R", + "--recursive", + action="store_true", + default=False, + help="List all stages inside the specified directory.", + ) + stage_list_parser.add_argument( + "--names-only", + action="store_true", + default=False, + help="List only the name of the stages.", + ) + stage_list_parser.set_defaults(func=CmdStageList) diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 95bda4cdb..186c17bf6 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -2,7 +2,7 @@ import logging import os from contextlib import contextmanager from functools import wraps -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Callable, Optional from funcy import cached_property, cat @@ -187,7 +187,9 @@ class Repo: self.params = Params(self) self.live = Live(self) - self.stage_collection_error_handler = None + self.stage_collection_error_handler: Optional[ + Callable[[str, Exception], None] + ] = None self._lock_depth = 0 @cached_property diff --git a/dvc/repo/stage.py b/dvc/repo/stage.py index 15df756ad..91f5217d8 100644 --- a/dvc/repo/stage.py +++ b/dvc/repo/stage.py @@ -3,7 +3,16 @@ import logging import os import typing from contextlib import suppress -from typing import Iterable, List, NamedTuple, Optional, Set, Tuple, Union +from typing import ( + Callable, + Iterable, + List, + NamedTuple, + Optional, + Set, + Tuple, + Union, +) from dvc.exceptions import ( DvcException, @@ -351,20 +360,20 @@ class StageLoad: return [StageInfo(stage) for stage in stages] - def collect_repo(self, onerror=None): + def collect_repo(self, onerror: Callable[[str, Exception], None] = None): """Collects all of the stages present in the DVC repo. Args: onerror (optional): callable that will be called with two args: the filepath whose collection failed and the exc instance. - The decision to raise error is with the callback. - Otherwise, the file will be skipped, and dvc will partially - load the stages. + It can report the error to continue with the collection + (and, skip failed ones), or raise the exception to abort + the collection. """ from dvc.dvcfile import is_valid_filename stages = [] - outs = set() + outs: Set[str] = set() for root, dirs, files in self.tree.walk(self.repo.root_dir): for file_name in filter(is_valid_filename, files): file_path = os.path.join(root, file_name) diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 990d3016d..e0854fec0 100644 --- a/dvc/stage/__init__.py +++ b/dvc/stage/__init__.py @@ -2,6 +2,7 @@ import logging import os import string from collections import defaultdict +from contextlib import suppress from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, Optional @@ -135,7 +136,7 @@ class Stage(params.StageParams): always_changed=False, stage_text=None, dvcfile=None, - desc=None, + desc: Optional[str] = None, meta=None, ): if deps is None: @@ -154,7 +155,7 @@ class Stage(params.StageParams): self.always_changed = always_changed self._stage_text = stage_text self._dvcfile = dvcfile - self.desc = desc + self.desc: Optional[str] = desc self.meta = meta self.raw_data = RawData() @@ -255,6 +256,15 @@ class Stage(params.StageParams): """ return any(out.checkpoint for out in self.outs) + def short_description(self) -> Optional["str"]: + desc: Optional["str"] = None + if self.desc: + with suppress(ValueError): + # try to use first non-empty line as a description + line = next(filter(None, self.desc.splitlines())) + desc = line.strip() + return desc + def _read_env(self, out, checkpoint_func=None) -> Env: env: Env = {} if out.live: