Move Repo.collect and collect_granular to Repo.stage (#4993)

dvc/command/dag.py

@@ -34,7 +34,7 @@ def _collect_targets(repo, target, outs):
     if not target:
         return []
 
-    pairs = repo.collect_granular(target)
+    pairs = repo.stage.collect_granular(target)
     if not outs:
         return [stage.addressing for stage, _ in pairs]
 

dvc/repo/__init__.py

@@ -1,37 +1,27 @@
 import logging
 import os
-from contextlib import contextmanager, suppress
+from contextlib import contextmanager
 from functools import wraps
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING
 
 from funcy import cached_property, cat
 from git import InvalidGitRepositoryError
 
 from dvc.config import Config
-from dvc.dvcfile import PIPELINE_FILE, is_valid_filename
+from dvc.dvcfile import is_valid_filename
 from dvc.exceptions import FileMissingError
 from dvc.exceptions import IsADirectoryError as DvcIsADirectoryError
-from dvc.exceptions import (
-    NoOutputOrStageError,
-    NotDvcRepoError,
-    OutputNotFoundError,
-)
+from dvc.exceptions import NotDvcRepoError, OutputNotFoundError
 from dvc.path_info import PathInfo
-from dvc.repo.stage import StageLoad
 from dvc.scm import Base
 from dvc.scm.base import SCMError
 from dvc.tree.repo import RepoTree
 from dvc.utils.fs import path_isin
 
-from ..stage.exceptions import StageFileDoesNotExistError, StageNotFound
-from ..utils import parse_target
-from .graph import build_graph, build_outs_graph, get_pipeline, get_pipelines
+from .graph import build_graph, build_outs_graph, get_pipelines
 from .trie import build_outs_trie
 
 if TYPE_CHECKING:
-    from networkx import DiGraph
-
-    from dvc.stage import Stage
     from dvc.tree.base import BaseTree
 
 
@@ -147,6 +137,7 @@ class Repo:
         from dvc.repo.metrics import Metrics
         from dvc.repo.params import Params
         from dvc.repo.plots import Plots
+        from dvc.repo.stage import StageLoad
         from dvc.stage.cache import StageCache
         from dvc.state import State, StateNoop
         from dvc.tree.local import LocalTree
@@ -292,132 +283,6 @@ class Repo:
         if not getattr(self, "_skip_graph_checks", False):
             build_graph(self.stages + new_stages)
 
-    @staticmethod
-    def _collect_inside(path: str, graph: "DiGraph"):
-        import networkx as nx
-
-        stages = nx.dfs_postorder_nodes(graph)
-        return [stage for stage in stages if path_isin(stage.path, path)]
-
-    def collect(
-        self,
-        target: str = None,
-        with_deps: bool = False,
-        recursive: bool = False,
-        graph: "DiGraph" = None,
-        accept_group: bool = False,
-        glob: bool = False,
-    ):
-        if not target:
-            return list(graph) if graph else self.stages
-
-        if recursive and self.tree.isdir(target):
-            return self._collect_inside(
-                os.path.abspath(target), graph or self.graph
-            )
-
-        stages = self.stage.from_target(
-            target, accept_group=accept_group, glob=glob
-        )
-        if not with_deps:
-            return stages
-
-        return self._collect_stages_with_deps(stages, graph=graph)
-
-    def _collect_stages_with_deps(
-        self, stages: List["Stage"], graph: "DiGraph" = None
-    ):
-        res = set()
-        for stage in stages:
-            res.update(self._collect_pipeline(stage, graph=graph))
-        return res
-
-    def _collect_pipeline(self, stage: "Stage", graph: "DiGraph" = None):
-        import networkx as nx
-
-        pipeline = get_pipeline(get_pipelines(graph or self.graph), stage)
-        return nx.dfs_postorder_nodes(pipeline, stage)
-
-    def _collect_specific_target(
-        self, target: str, with_deps: bool, recursive: bool, accept_group: bool
-    ):
-        # Optimization: do not collect the graph for a specific target
-        file, name = parse_target(target)
-        stages = []
-        if not file:
-            # parsing is ambiguous when it does not have a colon
-            # or if it's not a dvcfile, as it can be a stage name
-            # in `dvc.yaml` or, an output in a stage.
-            logger.debug(
-                "Checking if stage '%s' is in '%s'", target, PIPELINE_FILE
-            )
-            if not (
-                recursive and self.tree.isdir(target)
-            ) and self.tree.exists(PIPELINE_FILE):
-                with suppress(StageNotFound):
-                    stages = self.stage.load_all(
-                        PIPELINE_FILE, target, accept_group=accept_group
-                    )
-                    if with_deps:
-                        stages = self._collect_stages_with_deps(stages)
-
-        elif not with_deps and is_valid_filename(file):
-            stages = self.stage.load_all(
-                file, name, accept_group=accept_group,
-            )
-
-        return stages, file, name
-
-    def collect_granular(
-        self,
-        target: str = None,
-        with_deps: bool = False,
-        recursive: bool = False,
-        graph: "DiGraph" = None,
-        accept_group: bool = False,
-    ):
-        """
-        Priority is in the order of following in case of ambiguity:
-            - .dvc file or .yaml file
-            - dir if recursive and directory exists
-            - stage_name
-            - output file
-        """
-        if not target:
-            return [(stage, None) for stage in self.stages]
-
-        stages, file, _ = self._collect_specific_target(
-            target, with_deps, recursive, accept_group
-        )
-        if not stages:
-            if not (recursive and self.tree.isdir(target)):
-                try:
-                    (out,) = self.find_outs_by_path(target, strict=False)
-                    filter_info = PathInfo(os.path.abspath(target))
-                    return [(out.stage, filter_info)]
-                except OutputNotFoundError:
-                    pass
-
-            try:
-                stages = self.collect(
-                    target,
-                    with_deps,
-                    recursive,
-                    graph,
-                    accept_group=accept_group,
-                )
-            except StageFileDoesNotExistError as exc:
-                # collect() might try to use `target` as a stage name
-                # and throw error that dvc.yaml does not exist, whereas it
-                # should say that both stage name and file does not exist.
-                if file and is_valid_filename(file):
-                    raise
-                raise NoOutputOrStageError(target, exc.file) from exc
-            except StageNotFound as exc:
-                raise NoOutputOrStageError(target, exc.file) from exc
-
-        return [(stage, None) for stage in stages]
-
     def used_cache(
         self,
         targets=None,
@@ -457,7 +322,7 @@ class Repo:
             targets = targets or [None]
 
             pairs = cat(
-                self.collect_granular(
+                self.stage.collect_granular(
                     target, recursive=recursive, with_deps=with_deps
                 )
                 for target in targets
@@ -516,9 +381,8 @@ class Repo:
 
         for root, dirs, files in self.tree.walk(self.root_dir):
             for file_name in filter(is_valid_filename, files):
-                new_stages = self.stage.load_file(
-                    os.path.join(root, file_name)
-                )
+                file_path = os.path.join(root, file_name)
+                new_stages = self.stage.load_file(file_path)
                 stages.extend(new_stages)
                 outs.update(
                     out.fspath

dvc/repo/checkout.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from typing import TYPE_CHECKING, Set
 
 from dvc.exceptions import (
     CheckoutError,
@@ -11,6 +12,10 @@ from dvc.utils import relpath
 
 from . import locked
 
+if TYPE_CHECKING:
+    from . import Repo
+    from .stage import StageInfo
+
 logger = logging.getLogger(__name__)
 
 
@@ -38,6 +43,34 @@ def get_all_files_numbers(pairs):
     )
 
 
+def _collect_pairs(
+    self: "Repo", targets, with_deps: bool, recursive: bool
+) -> Set["StageInfo"]:
+    from dvc.stage.exceptions import (
+        StageFileBadNameError,
+        StageFileDoesNotExistError,
+    )
+
+    pairs: Set["StageInfo"] = set()
+    for target in targets:
+        try:
+            pairs.update(
+                self.stage.collect_granular(
+                    target, with_deps=with_deps, recursive=recursive
+                )
+            )
+        except (
+            StageFileDoesNotExistError,
+            StageFileBadNameError,
+            NoOutputOrStageError,
+        ) as exc:
+            if not target:
+                raise
+            raise CheckoutErrorSuggestGit(target) from exc
+
+    return pairs
+
+
 @locked
 def checkout(
     self,
@@ -49,10 +82,6 @@ def checkout(
     allow_missing=False,
     **kwargs,
 ):
-    from dvc.stage.exceptions import (
-        StageFileBadNameError,
-        StageFileDoesNotExistError,
-    )
 
     unused = []
     stats = {
@@ -71,23 +100,7 @@ def checkout(
     if isinstance(targets, str):
         targets = [targets]
 
-    pairs = set()
-    for target in targets:
-        try:
-            pairs.update(
-                self.collect_granular(
-                    target, with_deps=with_deps, recursive=recursive
-                )
-            )
-        except (
-            StageFileDoesNotExistError,
-            StageFileBadNameError,
-            NoOutputOrStageError,
-        ) as exc:
-            if not target:
-                raise
-            raise CheckoutErrorSuggestGit(target) from exc
-
+    pairs = _collect_pairs(self, targets, with_deps, recursive)
     total = get_all_files_numbers(pairs)
     with Tqdm(
         total=total, unit="file", desc="Checkout", disable=total == 0

dvc/repo/commit.py

@@ -36,8 +36,12 @@ def prompt_to_commit(stage, changes, force=False):
 
 
 @locked
-def commit(self, target, with_deps=False, recursive=False, force=False):
-    stages = self.collect(target, with_deps=with_deps, recursive=recursive)
+def commit(
+    self, target, with_deps=False, recursive=False, force=False,
+):
+    stages = self.stage.collect(
+        target, with_deps=with_deps, recursive=recursive
+    )
     for stage in stages:
         changes = stage.changed_entries()
         if any(changes):

dvc/repo/graph.py

@@ -1,4 +1,14 @@
-def check_acyclic(graph):
+from typing import TYPE_CHECKING, Iterator, List, Set
+
+from dvc.utils.fs import path_isin
+
+if TYPE_CHECKING:
+    from networkx import DiGraph
+
+    from dvc.stage import Stage
+
+
+def check_acyclic(graph: "DiGraph") -> None:
     import networkx as nx
 
     from dvc.exceptions import CyclicGraphError
@@ -8,7 +18,7 @@ def check_acyclic(graph):
     except nx.NetworkXNoCycle:
         return
 
-    stages = set()
+    stages: Set["Stage"] = set()
     for from_node, to_node, _ in edges:
         stages.add(from_node)
         stages.add(to_node)
@@ -28,6 +38,20 @@ def get_pipelines(G):
     return [G.subgraph(c).copy() for c in nx.weakly_connected_components(G)]
 
 
+def collect_pipeline(stage: "Stage", graph: "DiGraph") -> Iterator["Stage"]:
+    import networkx as nx
+
+    pipeline = get_pipeline(get_pipelines(graph), stage)
+    return nx.dfs_postorder_nodes(pipeline, stage)
+
+
+def collect_inside_path(path: str, graph: "DiGraph") -> List["Stage"]:
+    import networkx as nx
+
+    stages = nx.dfs_postorder_nodes(graph)
+    return [stage for stage in stages if path_isin(stage.path, path)]
+
+
 def build_graph(stages, outs_trie=None):
     """Generate a graph by using the given stages on the given directory
 

dvc/repo/reproduce.py

@@ -114,7 +114,7 @@ def reproduce(
                 if pipeline.in_degree(stage) == 0:
                     targets.append(stage)
     else:
-        targets = self.collect(
+        targets = self.stage.collect(
             target,
             recursive=recursive,
             graph=active_graph,

dvc/repo/stage.py

@@ -1,17 +1,84 @@
 import fnmatch
 import logging
+import os
 import typing
-from typing import Iterable, List
+from contextlib import suppress
+from typing import Iterable, List, NamedTuple, Optional, Set, Tuple
 
-from dvc.dvcfile import PIPELINE_FILE, Dvcfile
+from dvc.dvcfile import PIPELINE_FILE, Dvcfile, is_valid_filename
+from dvc.exceptions import NoOutputOrStageError, OutputNotFoundError
+from dvc.path_info import PathInfo
+from dvc.repo.graph import collect_inside_path, collect_pipeline
+from dvc.stage.exceptions import StageFileDoesNotExistError, StageNotFound
 from dvc.utils import parse_target
 
 logger = logging.getLogger(__name__)
 
 if typing.TYPE_CHECKING:
+    from networkx import DiGraph
+
     from dvc.repo import Repo
     from dvc.stage import Stage
     from dvc.stage.loader import StageLoader
+    from dvc.types import OptStr
+
+
+class StageInfo(NamedTuple):
+    stage: "Stage"
+    filter_info: Optional[PathInfo] = None
+
+
+StageList = List["Stage"]
+StageIter = Iterable["Stage"]
+StageSet = Set["Stage"]
+
+
+def _collect_with_deps(stages: StageList, graph: "DiGraph") -> StageSet:
+    res: StageSet = set()
+    for stage in stages:
+        res.update(collect_pipeline(stage, graph=graph))
+    return res
+
+
+def _maybe_collect_from_dvc_yaml(
+    loader: "StageLoad", target, with_deps: bool, **load_kwargs,
+) -> StageIter:
+
+    stages: StageList = []
+    if loader.tree.exists(PIPELINE_FILE):
+        with suppress(StageNotFound):
+            stages = loader.load_all(PIPELINE_FILE, target, **load_kwargs)
+    return _collect_with_deps(stages, loader.graph) if with_deps else stages
+
+
+def _collect_specific_target(
+    loader: "StageLoad",
+    target: str,
+    with_deps: bool,
+    recursive: bool,
+    accept_group: bool,
+) -> Tuple[StageIter, "OptStr", "OptStr"]:
+    # Optimization: do not collect the graph for a specific target
+    file, name = parse_target(target)
+
+    # if the target has a file, we can load directly from it.
+    if not file:
+        # but, if there's no file, parsing is ambiguous as it can be a
+        # stage name in `dvc.yaml` file or an output. We prioritize
+        # `dvc.yaml` stage name here. If it exists, then we move on.
+        # else, we assume it's a output name in the `collect_granular()` below
+        msg = "Checking if stage '%s' is in '%s'"
+        logger.debug(msg, target, PIPELINE_FILE)
+        if not (recursive and loader.tree.isdir(target)):
+            stages = _maybe_collect_from_dvc_yaml(
+                loader, target, with_deps, accept_group=accept_group,
+            )
+            if stages:
+                return stages, file, name
+    elif not with_deps and is_valid_filename(file):
+        stages = loader.load_all(file, name, accept_group=accept_group)
+        return stages, file, name
+    return [], file, name
 
 
 class StageLoad:
@@ -20,7 +87,7 @@ class StageLoad:
 
     def from_target(
         self, target: str, accept_group: bool = False, glob: bool = False,
-    ) -> List["Stage"]:
+    ) -> StageList:
         """
         Returns a list of stage from the provided target.
         (see load method below for further details)
@@ -81,7 +148,7 @@ class StageLoad:
         name: str = None,
         accept_group: bool = False,
         glob: bool = False,
-    ) -> List["Stage"]:
+    ) -> StageList:
         """Load a list of stages from a file.
 
         Args:
@@ -121,10 +188,131 @@ class StageLoad:
 
         return stages[name]
 
-    def load_file(self, path: str = None) -> List["Stage"]:
+    def load_file(self, path: str = None) -> StageList:
         """Load all of the stages from a file."""
         return self.load_all(path)
 
     def load_glob(self, path: str, expr: str = None):
         """Load stages from `path`, filtered with `expr` provided."""
         return self.load_all(path, expr, glob=True)
+
+    @property
+    def tree(self):
+        return self.repo.tree
+
+    @property
+    def graph(self) -> "DiGraph":
+        return self.repo.graph
+
+    def collect(
+        self,
+        target: str = None,
+        with_deps: bool = False,
+        recursive: bool = False,
+        graph: "DiGraph" = None,
+        accept_group: bool = False,
+        glob: bool = False,
+    ) -> StageIter:
+        """Collect list of stages from the provided target.
+
+        Args:
+            target: if not provided, all of the stages in the graph are
+                returned.
+                Target can be:
+                - a stage name in the `dvc.yaml` file.
+                - a path to `dvc.yaml` or `.dvc` file.
+                - in case of a stage to a dvc.yaml file in a different
+                  directory than current working directory, it can be a path
+                  to dvc.yaml file, followed by a colon `:`, followed by stage
+                  name (eg: `../dvc.yaml:build`).
+                - in case of `recursive`, it can be a path to a directory.
+                - in case of `accept_group`, it can be a group name of
+                    `foreach` generated stage.
+                - in case of `glob`, it can be a wildcard pattern to match
+                  stages. Example: `build*` for stages in `dvc.yaml` file, or
+                  `../dvc.yaml:build*` for stages in dvc.yaml in a different
+                  directory.
+                  Note that, glob only applies for the stage name, not to the
+                  file, so `**/dvc.yaml:build*` is not possible.
+            with_deps: if true, the stages including their dependencies are
+                returned.
+            recursive: if true and if `target` is a directory, all of the
+                stages inside that directory is returned.
+            graph: graph to use. Defaults to `repo.graph`.
+            accept_group: if true, all of the `foreach` generated stages of
+                the specified target is returned.
+            glob: Use `target` as a pattern to match stages in a file.
+        """
+        if not target:
+            return list(graph) if graph else self.repo.stages
+
+        if recursive and self.repo.tree.isdir(target):
+            path = os.path.abspath(target)
+            return collect_inside_path(path, graph or self.graph)
+
+        stages = self.from_target(target, accept_group=accept_group, glob=glob)
+        if not with_deps:
+            return stages
+
+        return _collect_with_deps(stages, graph or self.graph)
+
+    def collect_granular(
+        self,
+        target: str = None,
+        with_deps: bool = False,
+        recursive: bool = False,
+        graph: "DiGraph" = None,
+        accept_group: bool = False,
+    ) -> List[StageInfo]:
+        """Collects a list of (stage, filter_info) from the given target.
+
+        Priority is in the order of following in case of ambiguity:
+        - .dvc file or .yaml file
+        - dir if recursive and directory exists
+        - stage_name
+        - output file
+
+        Args:
+            target: if not provided, all of the stages without any filters are
+                returned.
+                If `target` is a path to a dvc-tracked output,
+                a (stage, output_path_info) is returned.
+                Otherwise, the details above for `target` in `collect()`
+                applies.
+
+            (see `collect()` for other arguments)
+        """
+        if not target:
+            return [StageInfo(stage) for stage in self.repo.stages]
+
+        stages, file, _ = _collect_specific_target(
+            self, target, with_deps, recursive, accept_group
+        )
+        if not stages:
+            if not (recursive and self.tree.isdir(target)):
+                try:
+                    (out,) = self.repo.find_outs_by_path(target, strict=False)
+                    filter_info = PathInfo(os.path.abspath(target))
+                    return [StageInfo(out.stage, filter_info)]
+                except OutputNotFoundError:
+                    pass
+
+            try:
+                stages = self.collect(
+                    target,
+                    with_deps,
+                    recursive,
+                    graph,
+                    accept_group=accept_group,
+                )
+            except StageFileDoesNotExistError as exc:
+                # collect() might try to use `target` as a stage name
+                # and throw error that dvc.yaml does not exist, whereas it
+                # should say that both stage name and file does not exist.
+                if file and is_valid_filename(file):
+                    raise
+                raise NoOutputOrStageError(target, exc.file) from exc
+            except StageNotFound as exc:
+                raise NoOutputOrStageError(target, exc.file) from exc
+
+        return [StageInfo(stage) for stage in stages]

dvc/repo/status.py

@@ -29,7 +29,9 @@ def _joint_status(pairs):
 def _local_status(self, targets=None, with_deps=False, recursive=False):
     targets = targets or [None]
     pairs = cat(
-        self.collect_granular(t, with_deps=with_deps, recursive=recursive)
+        self.stage.collect_granular(
+            t, with_deps=with_deps, recursive=recursive
+        )
         for t in targets
     )
 

dvc/repo/update.py

@@ -3,7 +3,7 @@ from . import locked
 
 
 @locked
-def update(self, targets, rev=None, recursive=False):
+def update(self, targets=None, rev=None, recursive=False):
     if not targets:
         targets = [None]
 
@@ -12,7 +12,7 @@ def update(self, targets, rev=None, recursive=False):
 
     stages = set()
     for target in targets:
-        stages.update(self.collect(target, recursive=recursive))
+        stages.update(self.stage.collect(target, recursive=recursive))
 
     for stage in stages:
         stage.update(rev)

dvc/types.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from dvc.path_info import PathInfo, URLInfo
 
@@ -8,3 +8,5 @@ if TYPE_CHECKING:
 StrPath = Union[str, "PathLike[str]"]
 DvcPath = Union[PathInfo, URLInfo]
 AnyPath = Union[str, DvcPath, StrPath]
+
+OptStr = Optional[str]

tests/func/test_repo.py

@@ -1,22 +1,6 @@
-import os
-from operator import itemgetter
-
-import pytest
-
 from dvc.cache import Cache
 from dvc.dvcfile import PIPELINE_FILE, PIPELINE_LOCK
-from dvc.exceptions import NoOutputOrStageError
-from dvc.path_info import PathInfo
-from dvc.repo import Repo
-from dvc.stage.exceptions import (
-    StageFileDoesNotExistError,
-    StageNameUnspecified,
-    StageNotFound,
-)
 from dvc.system import System
-from dvc.utils import relpath
-from dvc.utils.fs import remove
-from dvc.utils.serialize import dump_yaml
 
 
 def test_destroy(tmp_dir, dvc, run_copy):
@@ -66,368 +50,3 @@ def test_destroy(tmp_dir, dvc, run_copy):
         "dir/subdir/file",
     ]:
         assert not System.is_symlink(tmp_dir / path)
-
-
-def test_collect(tmp_dir, scm, dvc, run_copy):
-    def collect_outs(*args, **kwargs):
-        return {
-            str(out)
-            for stage in dvc.collect(*args, **kwargs)
-            for out in stage.outs
-        }
-
-    tmp_dir.dvc_gen("foo", "foo")
-    run_copy("foo", "bar", single_stage=True)
-    scm.add([".gitignore", "foo.dvc", "bar.dvc"])
-    scm.commit("Add foo and bar")
-
-    scm.checkout("new-branch", create_new=True)
-
-    run_copy("bar", "buzz", single_stage=True)
-    scm.add([".gitignore", "buzz.dvc"])
-    scm.commit("Add buzz")
-
-    assert collect_outs("bar.dvc", with_deps=True) == {"foo", "bar"}
-    assert collect_outs("buzz.dvc", with_deps=True) == {"foo", "bar", "buzz"}
-    assert collect_outs("buzz.dvc", with_deps=False) == {"buzz"}
-
-    run_copy("foo", "foobar", name="copy-foo-foobar")
-    assert collect_outs(":copy-foo-foobar") == {"foobar"}
-    assert collect_outs(":copy-foo-foobar", with_deps=True) == {
-        "foobar",
-        "foo",
-    }
-    assert collect_outs("dvc.yaml:copy-foo-foobar", recursive=True) == {
-        "foobar"
-    }
-    assert collect_outs("copy-foo-foobar") == {"foobar"}
-    assert collect_outs("copy-foo-foobar", with_deps=True) == {
-        "foobar",
-        "foo",
-    }
-    assert collect_outs("copy-foo-foobar", recursive=True) == {"foobar"}
-
-    run_copy("foobar", "baz", name="copy-foobar-baz")
-    assert collect_outs("dvc.yaml") == {"foobar", "baz"}
-    assert collect_outs("dvc.yaml", with_deps=True) == {
-        "foobar",
-        "baz",
-        "foo",
-    }
-
-
-def test_collect_dir_recursive(tmp_dir, dvc, run_head):
-    tmp_dir.gen({"dir": {"foo": "foo"}})
-    (stage1,) = dvc.add("dir", recursive=True)
-    with (tmp_dir / "dir").chdir():
-        stage2 = run_head("foo", name="copy-foo-bar")
-        stage3 = run_head("foo-1", single_stage=True)
-    assert set(dvc.collect("dir", recursive=True)) == {stage1, stage2, stage3}
-
-
-def test_collect_with_not_existing_output_or_stage_name(
-    tmp_dir, dvc, run_copy
-):
-    with pytest.raises(StageFileDoesNotExistError):
-        dvc.collect("some_file")
-    tmp_dir.dvc_gen("foo", "foo")
-    run_copy("foo", "bar", name="copy-foo-bar")
-    with pytest.raises(StageNotFound):
-        dvc.collect("some_file")
-
-
-def test_stages(tmp_dir, dvc):
-    def collect_stages():
-        return {stage.relpath for stage in Repo(os.fspath(tmp_dir)).stages}
-
-    tmp_dir.dvc_gen({"file": "a", "dir/file": "b", "dir/subdir/file": "c"})
-
-    assert collect_stages() == {
-        "file.dvc",
-        os.path.join("dir", "file.dvc"),
-        os.path.join("dir", "subdir", "file.dvc"),
-    }
-
-    tmp_dir.gen(".dvcignore", "dir")
-
-    assert collect_stages() == {"file.dvc"}
-
-
-@pytest.fixture
-def stages(tmp_dir, run_copy):
-    stage1, stage2 = tmp_dir.dvc_gen({"foo": "foo", "lorem": "lorem"})
-    return {
-        "foo-generate": stage1,
-        "lorem-generate": stage2,
-        "copy-foo-bar": run_copy("foo", "bar", single_stage=True),
-        "copy-bar-foobar": run_copy("bar", "foobar", name="copy-bar-foobar"),
-        "copy-lorem-ipsum": run_copy(
-            "lorem", "ipsum", name="copy-lorem-ipsum"
-        ),
-    }
-
-
-def test_collect_not_a_group_stage_with_group_flag(tmp_dir, dvc, stages):
-    assert set(dvc.collect("copy-bar-foobar", accept_group=True)) == {
-        stages["copy-bar-foobar"]
-    }
-    assert set(
-        dvc.collect("copy-bar-foobar", accept_group=True, with_deps=True)
-    ) == {
-        stages["copy-bar-foobar"],
-        stages["copy-foo-bar"],
-        stages["foo-generate"],
-    }
-    assert set(dvc.collect_granular("copy-bar-foobar", accept_group=True)) == {
-        (stages["copy-bar-foobar"], None),
-    }
-    assert set(
-        dvc.collect_granular(
-            "copy-bar-foobar", accept_group=True, with_deps=True
-        )
-    ) == {
-        (stages["copy-bar-foobar"], None),
-        (stages["copy-foo-bar"], None),
-        (stages["foo-generate"], None),
-    }
-
-
-def test_collect_generated(tmp_dir, dvc):
-    dvc.config["feature"]["parametrization"] = True
-    d = {
-        "vars": [{"vars": [1, 2, 3, 4, 5]}],
-        "stages": {
-            "build": {"foreach": "${vars}", "do": {"cmd": "echo ${item}"}}
-        },
-    }
-    dump_yaml("dvc.yaml", d)
-
-    all_stages = set(dvc.stages)
-    assert len(all_stages) == 5
-
-    assert set(dvc.collect()) == all_stages
-    assert set(dvc.collect("build", accept_group=True)) == all_stages
-    assert (
-        set(dvc.collect("build", accept_group=True, with_deps=True))
-        == all_stages
-    )
-    assert set(dvc.collect("build*", glob=True)) == all_stages
-    assert set(dvc.collect("build*", glob=True, with_deps=True)) == all_stages
-
-    stages_info = {(stage, None) for stage in all_stages}
-    assert set(dvc.collect_granular("build", accept_group=True)) == stages_info
-    assert (
-        set(dvc.collect_granular("build", accept_group=True, with_deps=True))
-        == stages_info
-    )
-
-
-def test_collect_glob(tmp_dir, dvc, stages):
-    assert set(dvc.collect("copy*", glob=True)) == {
-        stages[key] for key in ["copy-bar-foobar", "copy-lorem-ipsum"]
-    }
-    assert set(dvc.collect("copy-lorem*", glob=True, with_deps=True)) == {
-        stages[key] for key in ["copy-lorem-ipsum", "lorem-generate"]
-    }
-
-
-def test_collect_granular_with_no_target(tmp_dir, dvc, stages):
-    assert set(map(itemgetter(0), dvc.collect_granular())) == set(
-        stages.values()
-    )
-    assert list(map(itemgetter(1), dvc.collect_granular())) == [None] * len(
-        stages
-    )
-
-
-def test_collect_granular_with_target(tmp_dir, dvc, stages):
-    assert dvc.collect_granular("bar.dvc") == [(stages["copy-foo-bar"], None)]
-    assert dvc.collect_granular(PIPELINE_FILE) == [
-        (stages["copy-bar-foobar"], None),
-        (stages["copy-lorem-ipsum"], None),
-    ]
-    assert dvc.collect_granular(":") == [
-        (stages["copy-bar-foobar"], None),
-        (stages["copy-lorem-ipsum"], None),
-    ]
-    assert dvc.collect_granular("copy-bar-foobar") == [
-        (stages["copy-bar-foobar"], None)
-    ]
-    assert dvc.collect_granular(":copy-bar-foobar") == [
-        (stages["copy-bar-foobar"], None)
-    ]
-    assert dvc.collect_granular("dvc.yaml:copy-bar-foobar") == [
-        (stages["copy-bar-foobar"], None)
-    ]
-
-    with (tmp_dir / dvc.DVC_DIR).chdir():
-        assert dvc.collect_granular(
-            relpath(tmp_dir / PIPELINE_FILE) + ":copy-bar-foobar"
-        ) == [(stages["copy-bar-foobar"], None)]
-
-    assert dvc.collect_granular("foobar") == [
-        (stages["copy-bar-foobar"], PathInfo(tmp_dir / "foobar"))
-    ]
-
-
-@pytest.mark.parametrize(
-    "target",
-    [
-        "not_existing.dvc",
-        "not_existing.dvc:stage_name",
-        "not_existing/dvc.yaml",
-        "not_existing/dvc.yaml:stage_name",
-    ],
-)
-def test_collect_with_not_existing_dvcfile(tmp_dir, dvc, target):
-    with pytest.raises(StageFileDoesNotExistError):
-        dvc.collect_granular(target)
-    with pytest.raises(StageFileDoesNotExistError):
-        dvc.collect(target)
-
-
-def test_collect_granular_with_not_existing_output_or_stage_name(tmp_dir, dvc):
-    with pytest.raises(NoOutputOrStageError):
-        dvc.collect_granular("some_file")
-    with pytest.raises(NoOutputOrStageError):
-        dvc.collect_granular("some_file", recursive=True)
-
-
-def test_collect_granular_with_deps(tmp_dir, dvc, stages):
-    assert set(
-        map(itemgetter(0), dvc.collect_granular("bar.dvc", with_deps=True))
-    ) == {stages["copy-foo-bar"], stages["foo-generate"]}
-    assert set(
-        map(
-            itemgetter(0),
-            dvc.collect_granular("copy-bar-foobar", with_deps=True),
-        )
-    ) == {
-        stages["copy-bar-foobar"],
-        stages["copy-foo-bar"],
-        stages["foo-generate"],
-    }
-    assert set(
-        map(
-            itemgetter(0), dvc.collect_granular(PIPELINE_FILE, with_deps=True),
-        )
-    ) == set(stages.values())
-
-
-def test_collect_granular_same_output_name_stage_name(tmp_dir, dvc, run_copy):
-    (stage1,) = tmp_dir.dvc_gen("foo", "foo")
-    (stage2,) = tmp_dir.dvc_gen("copy-foo-bar", "copy-foo-bar")
-    stage3 = run_copy("foo", "bar", name="copy-foo-bar")
-
-    assert dvc.collect_granular("copy-foo-bar") == [(stage3, None)]
-
-    coll = dvc.collect_granular("copy-foo-bar", with_deps=True)
-    assert set(map(itemgetter(0), coll)) == {stage3, stage1}
-    assert list(map(itemgetter(1), coll)) == [None] * 2
-
-    assert dvc.collect_granular("./copy-foo-bar") == [
-        (stage2, PathInfo(tmp_dir / "copy-foo-bar"))
-    ]
-    assert dvc.collect_granular("./copy-foo-bar", with_deps=True) == [
-        (stage2, PathInfo(tmp_dir / "copy-foo-bar"))
-    ]
-
-
-def test_collect_granular_priority_on_collision(tmp_dir, dvc, run_copy):
-    tmp_dir.gen({"dir": {"foo": "foo"}, "foo": "foo"})
-    (stage1,) = dvc.add("dir", recursive=True)
-    stage2 = run_copy("foo", "bar", name="dir")
-
-    assert dvc.collect_granular("dir") == [(stage2, None)]
-    assert dvc.collect_granular("dir", recursive=True) == [(stage1, None)]
-
-    remove(tmp_dir / "dir")
-
-    assert dvc.collect_granular("dir") == [(stage2, None)]
-    assert dvc.collect_granular("dir", recursive=True) == [(stage2, None)]
-
-
-def test_collect_granular_collision_output_dir_stage_name(
-    tmp_dir, dvc, run_copy
-):
-    stage1, *_ = tmp_dir.dvc_gen({"dir": {"foo": "foo"}, "foo": "foo"})
-    stage3 = run_copy("foo", "bar", name="dir")
-
-    assert dvc.collect_granular("dir") == [(stage3, None)]
-    assert not dvc.collect_granular("dir", recursive=True)
-    assert dvc.collect_granular("./dir") == [
-        (stage1, PathInfo(tmp_dir / "dir"))
-    ]
-
-
-def test_collect_granular_not_existing_stage_name(tmp_dir, dvc, run_copy):
-    tmp_dir.dvc_gen("foo", "foo")
-    (stage,) = tmp_dir.dvc_gen("copy-foo-bar", "copy-foo-bar")
-    run_copy("foo", "bar", name="copy-foo-bar")
-
-    assert dvc.collect_granular("copy-foo-bar.dvc:stage_name_not_needed") == [
-        (stage, None)
-    ]
-    with pytest.raises(StageNotFound):
-        dvc.collect_granular("dvc.yaml:does-not-exist")
-
-
-def test_get_stages(tmp_dir, dvc, run_copy):
-    with pytest.raises(StageFileDoesNotExistError):
-        dvc.stage.load_all()
-
-    tmp_dir.gen("foo", "foo")
-    stage1 = run_copy("foo", "bar", name="copy-foo-bar")
-    stage2 = run_copy("bar", "foobar", name="copy-bar-foobar")
-
-    assert set(dvc.stage.load_all()) == {stage1, stage2}
-    assert set(dvc.stage.load_all(path=PIPELINE_FILE)) == {stage1, stage2}
-    assert set(dvc.stage.load_all(name="copy-bar-foobar")) == {stage2}
-    assert set(
-        dvc.stage.load_all(path=PIPELINE_FILE, name="copy-bar-foobar")
-    ) == {stage2}
-
-    with pytest.raises(StageFileDoesNotExistError):
-        dvc.stage.load_all(path=relpath(tmp_dir / ".." / PIPELINE_FILE))
-
-    with pytest.raises(StageNotFound):
-        dvc.stage.load_all(path=PIPELINE_FILE, name="copy")
-
-
-def test_get_stages_old_dvcfile(tmp_dir, dvc):
-    (stage1,) = tmp_dir.dvc_gen("foo", "foo")
-    assert set(dvc.stage.load_all("foo.dvc")) == {stage1}
-    assert set(dvc.stage.load_all("foo.dvc", name="foo-generate")) == {stage1}
-
-    with pytest.raises(StageFileDoesNotExistError):
-        dvc.stage.load_all(path=relpath(tmp_dir / ".." / "foo.dvc"))
-
-
-def test_get_stage(tmp_dir, dvc, run_copy):
-    tmp_dir.gen("foo", "foo")
-    stage1 = run_copy("foo", "bar", name="copy-foo-bar")
-
-    with pytest.raises(StageNameUnspecified):
-        dvc.stage.load_one()
-
-    with pytest.raises(StageNameUnspecified):
-        dvc.stage.load_one(path=PIPELINE_FILE)
-
-    assert (
-        dvc.stage.load_one(path=PIPELINE_FILE, name="copy-foo-bar") == stage1
-    )
-    assert dvc.stage.load_one(name="copy-foo-bar") == stage1
-
-    with pytest.raises(StageFileDoesNotExistError):
-        dvc.stage.load_one(path="something.yaml", name="name")
-
-    with pytest.raises(StageNotFound):
-        dvc.stage.load_one(name="random_name")
-
-
-def test_get_stage_single_stage_dvcfile(tmp_dir, dvc):
-    (stage1,) = tmp_dir.dvc_gen("foo", "foo")
-    assert dvc.stage.load_one("foo.dvc") == stage1
-    assert dvc.stage.load_one("foo.dvc", name="jpt") == stage1
-    with pytest.raises(StageFileDoesNotExistError):
-        dvc.stage.load_one(path="bar.dvc", name="name")

tests/func/test_stage.py

@@ -216,9 +216,9 @@ def test_parent_repo_collect_stages(tmp_dir, scm, dvc):
         deep_subrepo_dir.gen("subrepo_file", "subrepo file content")
         deep_subrepo.add("subrepo_file")
 
-    stages = dvc.collect(None)
-    subrepo_stages = subrepo.collect(None)
-    deep_subrepo_stages = deep_subrepo.collect(None)
+    stages = dvc.stage.collect(None)
+    subrepo_stages = subrepo.stage.collect(None)
+    deep_subrepo_stages = deep_subrepo.stage.collect(None)
 
     assert stages == []
     assert subrepo_stages != []

tests/func/test_stage_load.py

@@ -0,0 +1,437 @@
+import os
+from operator import itemgetter
+
+import pytest
+from funcy import raiser
+
+from dvc.dvcfile import PIPELINE_FILE
+from dvc.exceptions import NoOutputOrStageError
+from dvc.path_info import PathInfo
+from dvc.repo import Repo
+from dvc.stage.exceptions import (
+    StageFileDoesNotExistError,
+    StageNameUnspecified,
+    StageNotFound,
+)
+from dvc.utils import relpath
+from dvc.utils.fs import remove
+from dvc.utils.serialize import dump_yaml
+
+
+def test_collect(tmp_dir, scm, dvc, run_copy):
+    def collect_outs(*args, **kwargs):
+        return {
+            str(out)
+            for stage in dvc.stage.collect(*args, **kwargs)
+            for out in stage.outs
+        }
+
+    tmp_dir.dvc_gen("foo", "foo")
+    run_copy("foo", "bar", single_stage=True)
+    scm.add([".gitignore", "foo.dvc", "bar.dvc"])
+    scm.commit("Add foo and bar")
+
+    scm.checkout("new-branch", create_new=True)
+
+    run_copy("bar", "buzz", single_stage=True)
+    scm.add([".gitignore", "buzz.dvc"])
+    scm.commit("Add buzz")
+
+    assert collect_outs("bar.dvc", with_deps=True) == {"foo", "bar"}
+    assert collect_outs("buzz.dvc", with_deps=True) == {"foo", "bar", "buzz"}
+    assert collect_outs("buzz.dvc", with_deps=False) == {"buzz"}
+
+    run_copy("foo", "foobar", name="copy-foo-foobar")
+    assert collect_outs(":copy-foo-foobar") == {"foobar"}
+    assert collect_outs(":copy-foo-foobar", with_deps=True) == {
+        "foobar",
+        "foo",
+    }
+    assert collect_outs("dvc.yaml:copy-foo-foobar", recursive=True) == {
+        "foobar"
+    }
+    assert collect_outs("copy-foo-foobar") == {"foobar"}
+    assert collect_outs("copy-foo-foobar", with_deps=True) == {
+        "foobar",
+        "foo",
+    }
+    assert collect_outs("copy-foo-foobar", recursive=True) == {"foobar"}
+
+    run_copy("foobar", "baz", name="copy-foobar-baz")
+    assert collect_outs("dvc.yaml") == {"foobar", "baz"}
+    assert collect_outs("dvc.yaml", with_deps=True) == {
+        "foobar",
+        "baz",
+        "foo",
+    }
+
+
+def test_collect_dir_recursive(tmp_dir, dvc, run_head):
+    tmp_dir.gen({"dir": {"foo": "foo"}})
+    (stage1,) = dvc.add("dir", recursive=True)
+    with (tmp_dir / "dir").chdir():
+        stage2 = run_head("foo", name="copy-foo-bar")
+        stage3 = run_head("foo-1", single_stage=True)
+    assert set(dvc.stage.collect("dir", recursive=True)) == {
+        stage1,
+        stage2,
+        stage3,
+    }
+
+
+def test_collect_with_not_existing_output_or_stage_name(
+    tmp_dir, dvc, run_copy
+):
+    with pytest.raises(StageFileDoesNotExistError):
+        dvc.stage.collect("some_file")
+    tmp_dir.dvc_gen("foo", "foo")
+    run_copy("foo", "bar", name="copy-foo-bar")
+    with pytest.raises(StageNotFound):
+        dvc.stage.collect("some_file")
+
+
+def test_stages(tmp_dir, dvc):
+    def collect_stages():
+        return {stage.relpath for stage in Repo(os.fspath(tmp_dir)).stages}
+
+    tmp_dir.dvc_gen({"file": "a", "dir/file": "b", "dir/subdir/file": "c"})
+
+    assert collect_stages() == {
+        "file.dvc",
+        os.path.join("dir", "file.dvc"),
+        os.path.join("dir", "subdir", "file.dvc"),
+    }
+
+    tmp_dir.gen(".dvcignore", "dir")
+
+    assert collect_stages() == {"file.dvc"}
+
+
+@pytest.fixture
+def stages(tmp_dir, run_copy):
+    stage1, stage2 = tmp_dir.dvc_gen({"foo": "foo", "lorem": "lorem"})
+    return {
+        "foo-generate": stage1,
+        "lorem-generate": stage2,
+        "copy-foo-bar": run_copy("foo", "bar", single_stage=True),
+        "copy-bar-foobar": run_copy("bar", "foobar", name="copy-bar-foobar"),
+        "copy-lorem-ipsum": run_copy(
+            "lorem", "ipsum", name="copy-lorem-ipsum"
+        ),
+    }
+
+
+def test_collect_not_a_group_stage_with_group_flag(tmp_dir, dvc, stages):
+    assert set(dvc.stage.collect("copy-bar-foobar", accept_group=True)) == {
+        stages["copy-bar-foobar"]
+    }
+    assert set(
+        dvc.stage.collect("copy-bar-foobar", accept_group=True, with_deps=True)
+    ) == {
+        stages["copy-bar-foobar"],
+        stages["copy-foo-bar"],
+        stages["foo-generate"],
+    }
+    assert set(
+        dvc.stage.collect_granular("copy-bar-foobar", accept_group=True)
+    ) == {(stages["copy-bar-foobar"], None)}
+    assert set(
+        dvc.stage.collect_granular(
+            "copy-bar-foobar", accept_group=True, with_deps=True
+        )
+    ) == {
+        (stages["copy-bar-foobar"], None),
+        (stages["copy-foo-bar"], None),
+        (stages["foo-generate"], None),
+    }
+
+
+def test_collect_generated(tmp_dir, dvc):
+    dvc.config["feature"]["parametrization"] = True
+    d = {
+        "vars": [{"vars": [1, 2, 3, 4, 5]}],
+        "stages": {
+            "build": {"foreach": "${vars}", "do": {"cmd": "echo ${item}"}}
+        },
+    }
+    dump_yaml("dvc.yaml", d)
+
+    all_stages = set(dvc.stages)
+    assert len(all_stages) == 5
+
+    assert set(dvc.stage.collect()) == all_stages
+    assert set(dvc.stage.collect("build", accept_group=True)) == all_stages
+    assert (
+        set(dvc.stage.collect("build", accept_group=True, with_deps=True))
+        == all_stages
+    )
+    assert set(dvc.stage.collect("build*", glob=True)) == all_stages
+    assert (
+        set(dvc.stage.collect("build*", glob=True, with_deps=True))
+        == all_stages
+    )
+
+    stages_info = {(stage, None) for stage in all_stages}
+    assert (
+        set(dvc.stage.collect_granular("build", accept_group=True))
+        == stages_info
+    )
+    assert (
+        set(
+            dvc.stage.collect_granular(
+                "build", accept_group=True, with_deps=True
+            )
+        )
+        == stages_info
+    )
+
+
+def test_collect_glob(tmp_dir, dvc, stages):
+    assert set(dvc.stage.collect("copy*", glob=True)) == {
+        stages[key] for key in ["copy-bar-foobar", "copy-lorem-ipsum"]
+    }
+    assert set(
+        dvc.stage.collect("copy-lorem*", glob=True, with_deps=True)
+    ) == {stages[key] for key in ["copy-lorem-ipsum", "lorem-generate"]}
+
+
+def test_collect_granular_with_no_target(tmp_dir, dvc, stages):
+    assert set(map(itemgetter(0), dvc.stage.collect_granular())) == set(
+        stages.values()
+    )
+    assert list(map(itemgetter(1), dvc.stage.collect_granular())) == [
+        None
+    ] * len(stages)
+
+
+def test_collect_granular_with_target(tmp_dir, dvc, stages):
+    assert dvc.stage.collect_granular("bar.dvc") == [
+        (stages["copy-foo-bar"], None)
+    ]
+    assert dvc.stage.collect_granular(PIPELINE_FILE) == [
+        (stages["copy-bar-foobar"], None),
+        (stages["copy-lorem-ipsum"], None),
+    ]
+    assert dvc.stage.collect_granular(":") == [
+        (stages["copy-bar-foobar"], None),
+        (stages["copy-lorem-ipsum"], None),
+    ]
+    assert dvc.stage.collect_granular("copy-bar-foobar") == [
+        (stages["copy-bar-foobar"], None)
+    ]
+    assert dvc.stage.collect_granular(":copy-bar-foobar") == [
+        (stages["copy-bar-foobar"], None)
+    ]
+    assert dvc.stage.collect_granular("dvc.yaml:copy-bar-foobar") == [
+        (stages["copy-bar-foobar"], None)
+    ]
+
+    with (tmp_dir / dvc.DVC_DIR).chdir():
+        assert dvc.stage.collect_granular(
+            relpath(tmp_dir / PIPELINE_FILE) + ":copy-bar-foobar"
+        ) == [(stages["copy-bar-foobar"], None)]
+
+    assert dvc.stage.collect_granular("foobar") == [
+        (stages["copy-bar-foobar"], PathInfo(tmp_dir / "foobar"))
+    ]
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        "not_existing.dvc",
+        "not_existing.dvc:stage_name",
+        "not_existing/dvc.yaml",
+        "not_existing/dvc.yaml:stage_name",
+    ],
+)
+def test_collect_with_not_existing_dvcfile(tmp_dir, dvc, target):
+    with pytest.raises(StageFileDoesNotExistError):
+        dvc.stage.collect_granular(target)
+    with pytest.raises(StageFileDoesNotExistError):
+        dvc.stage.collect(target)
+
+
+def test_collect_granular_with_not_existing_output_or_stage_name(tmp_dir, dvc):
+    with pytest.raises(NoOutputOrStageError):
+        dvc.stage.collect_granular("some_file")
+    with pytest.raises(NoOutputOrStageError):
+        dvc.stage.collect_granular("some_file", recursive=True)
+
+
+def test_collect_granular_with_deps(tmp_dir, dvc, stages):
+    assert set(
+        map(
+            itemgetter(0),
+            dvc.stage.collect_granular("bar.dvc", with_deps=True),
+        )
+    ) == {stages["copy-foo-bar"], stages["foo-generate"]}
+    assert set(
+        map(
+            itemgetter(0),
+            dvc.stage.collect_granular("copy-bar-foobar", with_deps=True),
+        )
+    ) == {
+        stages["copy-bar-foobar"],
+        stages["copy-foo-bar"],
+        stages["foo-generate"],
+    }
+    assert set(
+        map(
+            itemgetter(0),
+            dvc.stage.collect_granular(PIPELINE_FILE, with_deps=True),
+        )
+    ) == set(stages.values())
+
+
+def test_collect_granular_same_output_name_stage_name(tmp_dir, dvc, run_copy):
+    (stage1,) = tmp_dir.dvc_gen("foo", "foo")
+    (stage2,) = tmp_dir.dvc_gen("copy-foo-bar", "copy-foo-bar")
+    stage3 = run_copy("foo", "bar", name="copy-foo-bar")
+
+    assert dvc.stage.collect_granular("copy-foo-bar") == [(stage3, None)]
+
+    coll = dvc.stage.collect_granular("copy-foo-bar", with_deps=True)
+    assert set(map(itemgetter(0), coll)) == {stage3, stage1}
+    assert list(map(itemgetter(1), coll)) == [None] * 2
+
+    assert dvc.stage.collect_granular("./copy-foo-bar") == [
+        (stage2, PathInfo(tmp_dir / "copy-foo-bar"))
+    ]
+    assert dvc.stage.collect_granular("./copy-foo-bar", with_deps=True) == [
+        (stage2, PathInfo(tmp_dir / "copy-foo-bar"))
+    ]
+
+
+def test_collect_granular_priority_on_collision(tmp_dir, dvc, run_copy):
+    tmp_dir.gen({"dir": {"foo": "foo"}, "foo": "foo"})
+    (stage1,) = dvc.add("dir", recursive=True)
+    stage2 = run_copy("foo", "bar", name="dir")
+
+    assert dvc.stage.collect_granular("dir") == [(stage2, None)]
+    assert dvc.stage.collect_granular("dir", recursive=True) == [
+        (stage1, None)
+    ]
+
+    remove(tmp_dir / "dir")
+
+    assert dvc.stage.collect_granular("dir") == [(stage2, None)]
+    assert dvc.stage.collect_granular("dir", recursive=True) == [
+        (stage2, None)
+    ]
+
+
+def test_collect_granular_collision_output_dir_stage_name(
+    tmp_dir, dvc, run_copy
+):
+    stage1, *_ = tmp_dir.dvc_gen({"dir": {"foo": "foo"}, "foo": "foo"})
+    stage3 = run_copy("foo", "bar", name="dir")
+
+    assert dvc.stage.collect_granular("dir") == [(stage3, None)]
+    assert not dvc.stage.collect_granular("dir", recursive=True)
+    assert dvc.stage.collect_granular("./dir") == [
+        (stage1, PathInfo(tmp_dir / "dir"))
+    ]
+
+
+def test_collect_granular_not_existing_stage_name(tmp_dir, dvc, run_copy):
+    tmp_dir.dvc_gen("foo", "foo")
+    (stage,) = tmp_dir.dvc_gen("copy-foo-bar", "copy-foo-bar")
+    run_copy("foo", "bar", name="copy-foo-bar")
+
+    assert dvc.stage.collect_granular(
+        "copy-foo-bar.dvc:stage_name_not_needed"
+    ) == [(stage, None)]
+    with pytest.raises(StageNotFound):
+        dvc.stage.collect_granular("dvc.yaml:does-not-exist")
+
+
+def test_get_stages(tmp_dir, dvc, run_copy):
+    with pytest.raises(StageFileDoesNotExistError):
+        dvc.stage.load_all()
+
+    tmp_dir.gen("foo", "foo")
+    stage1 = run_copy("foo", "bar", name="copy-foo-bar")
+    stage2 = run_copy("bar", "foobar", name="copy-bar-foobar")
+
+    assert set(dvc.stage.load_all()) == {stage1, stage2}
+    assert set(dvc.stage.load_all(path=PIPELINE_FILE)) == {stage1, stage2}
+    assert set(dvc.stage.load_all(name="copy-bar-foobar")) == {stage2}
+    assert set(
+        dvc.stage.load_all(path=PIPELINE_FILE, name="copy-bar-foobar")
+    ) == {stage2}
+
+    with pytest.raises(StageFileDoesNotExistError):
+        dvc.stage.load_all(path=relpath(tmp_dir / ".." / PIPELINE_FILE))
+
+    with pytest.raises(StageNotFound):
+        dvc.stage.load_all(path=PIPELINE_FILE, name="copy")
+
+
+def test_get_stages_old_dvcfile(tmp_dir, dvc):
+    (stage1,) = tmp_dir.dvc_gen("foo", "foo")
+    assert set(dvc.stage.load_all("foo.dvc")) == {stage1}
+    assert set(dvc.stage.load_all("foo.dvc", name="foo-generate")) == {stage1}
+
+    with pytest.raises(StageFileDoesNotExistError):
+        dvc.stage.load_all(path=relpath(tmp_dir / ".." / "foo.dvc"))
+
+
+def test_get_stage(tmp_dir, dvc, run_copy):
+    tmp_dir.gen("foo", "foo")
+    stage1 = run_copy("foo", "bar", name="copy-foo-bar")
+
+    with pytest.raises(StageNameUnspecified):
+        dvc.stage.load_one()
+
+    with pytest.raises(StageNameUnspecified):
+        dvc.stage.load_one(path=PIPELINE_FILE)
+
+    assert (
+        dvc.stage.load_one(path=PIPELINE_FILE, name="copy-foo-bar") == stage1
+    )
+    assert dvc.stage.load_one(name="copy-foo-bar") == stage1
+
+    with pytest.raises(StageFileDoesNotExistError):
+        dvc.stage.load_one(path="something.yaml", name="name")
+
+    with pytest.raises(StageNotFound):
+        dvc.stage.load_one(name="random_name")
+
+
+def test_get_stage_single_stage_dvcfile(tmp_dir, dvc):
+    (stage1,) = tmp_dir.dvc_gen("foo", "foo")
+    assert dvc.stage.load_one("foo.dvc") == stage1
+    assert dvc.stage.load_one("foo.dvc", name="jpt") == stage1
+    with pytest.raises(StageFileDoesNotExistError):
+        dvc.stage.load_one(path="bar.dvc", name="name")
+
+
+def test_collect_optimization(tmp_dir, dvc, mocker):
+    (stage,) = tmp_dir.dvc_gen("foo", "foo text")
+
+    # Forget cached stages and graph and error out on collection
+    dvc._reset()
+    mocker.patch(
+        "dvc.repo.Repo.stages",
+        property(raiser(Exception("Should not collect"))),
+    )
+
+    # Should read stage directly instead of collecting the whole graph
+    dvc.stage.collect(stage.path)
+    dvc.stage.collect_granular(stage.path)
+
+
+def test_collect_optimization_on_stage_name(tmp_dir, dvc, mocker, run_copy):
+    tmp_dir.dvc_gen("foo", "foo")
+    stage = run_copy("foo", "bar", name="copy-foo-bar")
+    # Forget cached stages and graph and error out on collection
+    dvc._reset()
+    mocker.patch(
+        "dvc.repo.Repo.stages",
+        property(raiser(Exception("Should not collect"))),
+    )
+
+    # Should read stage directly instead of collecting the whole graph
+    assert dvc.stage.collect("copy-foo-bar") == [stage]
+    assert dvc.stage.collect_granular("copy-foo-bar") == [(stage, None)]

tests/unit/repo/test_repo.py

@@ -1,7 +1,6 @@
 import os
 
 import pytest
-from funcy import raiser
 
 from dvc.repo import NotDvcRepoError, Repo, locked
 from dvc.utils.fs import remove
@@ -73,36 +72,6 @@ def test_locked(mocker):
     ]
 
 
-def test_collect_optimization(tmp_dir, dvc, mocker):
-    (stage,) = tmp_dir.dvc_gen("foo", "foo text")
-
-    # Forget cached stages and graph and error out on collection
-    dvc._reset()
-    mocker.patch(
-        "dvc.repo.Repo.stages",
-        property(raiser(Exception("Should not collect"))),
-    )
-
-    # Should read stage directly instead of collecting the whole graph
-    dvc.collect(stage.path)
-    dvc.collect_granular(stage.path)
-
-
-def test_collect_optimization_on_stage_name(tmp_dir, dvc, mocker, run_copy):
-    tmp_dir.dvc_gen("foo", "foo")
-    stage = run_copy("foo", "bar", name="copy-foo-bar")
-    # Forget cached stages and graph and error out on collection
-    dvc._reset()
-    mocker.patch(
-        "dvc.repo.Repo.stages",
-        property(raiser(Exception("Should not collect"))),
-    )
-
-    # Should read stage directly instead of collecting the whole graph
-    assert dvc.collect("copy-foo-bar") == [stage]
-    assert dvc.collect_granular("copy-foo-bar") == [(stage, None)]
-
-
 def test_skip_graph_checks(tmp_dir, dvc, mocker, run_copy):
     # See https://github.com/iterative/dvc/issues/2671 for more info
     mock_build_graph = mocker.patch("dvc.repo.build_graph")