Introduce stage:list command with autocompletions (#5282)

* Implement stage:list

* change targets help message

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

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",

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)

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

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)

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: