diff --git a/dvc/command/dag.py b/dvc/command/dag.py index 442e0ff08..c166a6745 100644 --- a/dvc/command/dag.py +++ b/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] diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index dc0af97e6..78dd4b7aa 100644 --- a/dvc/repo/__init__.py +++ b/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 diff --git a/dvc/repo/checkout.py b/dvc/repo/checkout.py index 6b2914031..b73f079cd 100644 --- a/dvc/repo/checkout.py +++ b/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 diff --git a/dvc/repo/commit.py b/dvc/repo/commit.py index 331162c67..ba8452129 100644 --- a/dvc/repo/commit.py +++ b/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): diff --git a/dvc/repo/graph.py b/dvc/repo/graph.py index 41142ae77..024e633b0 100644 --- a/dvc/repo/graph.py +++ b/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 diff --git a/dvc/repo/reproduce.py b/dvc/repo/reproduce.py index ba21ccb46..99cdd4058 100644 --- a/dvc/repo/reproduce.py +++ b/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, diff --git a/dvc/repo/stage.py b/dvc/repo/stage.py index c9e07bfcd..a5043eb04 100644 --- a/dvc/repo/stage.py +++ b/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] diff --git a/dvc/repo/status.py b/dvc/repo/status.py index accf18609..fb4f86b01 100644 --- a/dvc/repo/status.py +++ b/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 ) diff --git a/dvc/repo/update.py b/dvc/repo/update.py index d0dac9a97..e28c0f650 100644 --- a/dvc/repo/update.py +++ b/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) diff --git a/dvc/types.py b/dvc/types.py index 8d426f943..f38b1727e 100644 --- a/dvc/types.py +++ b/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] diff --git a/tests/func/test_repo.py b/tests/func/test_repo.py index c768d759a..ae5d35f5b 100644 --- a/tests/func/test_repo.py +++ b/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") diff --git a/tests/func/test_stage.py b/tests/func/test_stage.py index df17bf440..28435c1ad 100644 --- a/tests/func/test_stage.py +++ b/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 != [] diff --git a/tests/func/test_stage_load.py b/tests/func/test_stage_load.py new file mode 100644 index 000000000..1f051d829 --- /dev/null +++ b/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)] diff --git a/tests/unit/repo/test_repo.py b/tests/unit/repo/test_repo.py index 44ae5d0d4..272d5b78e 100644 --- a/tests/unit/repo/test_repo.py +++ b/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")