Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-02 07:18:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Saugat Pachhai
2020-11-30 14:28:16 +05:45
committed by GitHub
parent 56322d0a53
commit bc1d10532a
14 changed files with 718 additions and 596 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/command/dag.py
+1 -1
View File
@@ -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
+9 -145
View File
@@ -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
+34 -21
View File
@@ -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
+6 -2
View File
@@ -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
+26 -2
View File
@@ -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
+1 -1
View File
@@ -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
+193 -5
View File
@@ -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
+3 -1
View File
@@ -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
+2 -2
View File
@@ -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
+3 -1
View File
@@ -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
-381
View File
@@ -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
+3 -3
View File
@@ -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
+437
View File
@@ -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
-31
View File
@@ -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")