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

index/stage: move repo collection logic to index

Browse Source
This commit is contained in:
Saugat Pachhai (सौगात)
2023-01-11 19:38:43 +02:00
committed by Ruslan Kuprieiev
parent a814f04ee0
commit 30ec1cf5e9
14 changed files with 308 additions and 453 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/dvcfile.py
+4 -12
View File
@@ -1,16 +1,7 @@
import contextlib
import logging
import os
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
List,
Tuple,
TypeVar,
Union,
)
from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeVar, Union
from funcy import cached_property
@@ -192,7 +183,7 @@ class SingleStageFile(FileMixin):
from dvc.stage.loader import SingleStageLoader as LOADER
metrics: List[str] = []
plots: Dict[str, Any] = {}
plots: Any = {}
params: List[str] = []
@property
@@ -242,6 +233,7 @@ class PipelineFile(FileMixin):
self.__dict__.pop("contents", None)
self.__dict__.pop("lockfile_contents", None)
self.__dict__.pop("resolver", None)
self.__dict__.pop("stages", None)
def dump(
self, stage, update_pipeline=True, update_lock=True, **kwargs
@@ -313,7 +305,7 @@ class PipelineFile(FileMixin):
wdir = self.repo.fs.path.parent(self.path)
return DataResolver(self.repo, wdir, self.contents)
@property
@cached_property
def stages(self):
return self.LOADER(self, self.contents, self.lockfile_contents)
dvc/repo/__init__.py
+8 -1
View File
@@ -239,7 +239,14 @@ class Repo:
def index(self):
from dvc.repo.index import Index
return Index(self)
return Index.from_repo(self)
def ensure_graph_correctness_with(self, stages=None, callback=None):
if not getattr(self, "_skip_graph_checks", False):
new = self.index.update(stages)
if callable(callback):
callback()
new.check_graph()
@staticmethod
def open(url, *args, **kwargs):
dvc/repo/add.py
+3 -3
View File
@@ -180,9 +180,9 @@ def add( # noqa: C901
with translate_graph_error(stages), ui.status(msg) as status:
# remove existing stages that are to-be replaced with these
# new stages for the graph checks.
new_index = repo.index.update(stages)
status.update("Checking graph")
new_index.check_graph()
repo.ensure_graph_correctness_with(
stages=stages, callback=lambda: status.update("Checking graph")
)
odb = None
if to_remote:
dvc/repo/imp_url.py
+1 -2
View File
@@ -78,8 +78,7 @@ def imp_url( # noqa: C901
dvcfile.remove()
try:
new_index = self.index.add(stage)
new_index.check_graph()
self.ensure_graph_correctness_with(stages={stage})
except OutputDuplicationError as exc:
raise OutputDuplicationError( # noqa: B904
exc.output, set(exc.stages) - {stage}
dvc/repo/index.py
+234 -195
View File
@@ -1,3 +1,6 @@
import logging
import time
from functools import partial
from typing import (
TYPE_CHECKING,
Any,
@@ -14,16 +17,18 @@ from typing import (
)
from funcy import cached_property
from funcy.debug import format_time
from dvc.fs import LocalFileSystem
if TYPE_CHECKING:
from networkx import DiGraph
from pygtrie import Trie
from dvc.dependency import Dependency, ParamsDependency
from dvc.fs import FileSystem
from dvc.output import Output
from dvc.repo import Repo
from dvc.repo.stage import StageInfo, StageLoad
from dvc.repo.stage import StageInfo
from dvc.stage import Stage
from dvc.types import TargetType
from dvc_data.hashfile.hash_info import HashInfo
@@ -31,174 +36,180 @@ if TYPE_CHECKING:
from dvc_objects.db import ObjectDB
logger = logging.getLogger(__name__)
ObjectContainer = Dict[Optional["ObjectDB"], Set["HashInfo"]]
def log_walk(seq):
for root, dirs, files in seq:
start = time.perf_counter()
yield root, dirs, files
duration = format_time(time.perf_counter() - start)
logger.trace("%s in collecting stages from %s", duration, root)
def collect_files(
repo: "Repo", onerror: Callable[[str, Exception], None] = None
):
"""Collects all of the stages present in the DVC repo.
Args:
onerror (optional): callable that will be called with two args:
the filepath whose collection failed and the exc instance.
It can report the error to continue with the collection
(and, skip failed ones), or raise the exception to abort
the collection.
"""
from dvc.dvcfile import is_valid_filename
from dvc.exceptions import DvcException
from dvc.utils import relpath
scm = repo.scm
fs = repo.fs
sep = fs.sep
outs: Set[str] = set()
is_local_fs = isinstance(fs, LocalFileSystem)
def is_ignored(path):
# apply only for the local fs
return is_local_fs and scm.is_ignored(path)
def is_dvcfile_and_not_ignored(root, file):
return is_valid_filename(file) and not is_ignored(f"{root}{sep}{file}")
def is_out_or_ignored(root, directory):
dir_path = f"{root}{sep}{directory}"
# trailing slash needed to check if a directory is gitignored
return dir_path in outs or is_ignored(f"{dir_path}{sep}")
walk_iter = repo.dvcignore.walk(fs, repo.root_dir, followlinks=False)
if logger.isEnabledFor(logging.TRACE): # type: ignore[attr-defined]
walk_iter = log_walk(walk_iter)
for root, dirs, files in walk_iter:
dvcfile_filter = partial(is_dvcfile_and_not_ignored, root)
for file in filter(dvcfile_filter, files):
file_path = fs.path.join(root, file)
try:
index = Index.from_file(repo, file_path)
except DvcException as exc:
if onerror:
onerror(relpath(file_path), exc)
continue
raise
outs.update(
out.fspath
for stage in index.stages
for out in stage.outs
if out.protocol == "local"
)
yield file_path, index
dirs[:] = [d for d in dirs if not is_out_or_ignored(root, d)]
class Index:
def __init__(
self,
repo: "Repo", # pylint: disable=redefined-outer-name
fs: "FileSystem" = None,
repo: "Repo",
stages: List["Stage"] = None,
metrics: Dict[str, List[str]] = None,
plots: Dict[str, List[str]] = None,
params: Dict[str, Any] = None,
) -> None:
"""Index is an immutable collection of stages.
Generally, Index is a complete collection of stages at a point in time.
With "a point in time", it means it is collected from the user's
workspace or a git revision.
And, since Index is immutable, the collection is frozen in time.
Index provides multiple ways to view this collection:
stages - provides direct access to this collection
outputs - provides direct access to the outputs
objects - provides direct access to the objects
graph -
... and many more.
Index also provides ways to slice and dice this collection.
Some `views` might not make sense when sliced (eg: pipelines/graph).
"""
self.repo: "Repo" = repo
self.fs: "FileSystem" = fs or repo.fs
self.stage_collector: "StageLoad" = repo.stage
if stages is not None:
self.stages: List["Stage"] = stages
self.repo = repo
self.stages = stages or []
self._metrics = metrics or {}
self._plots = plots or {}
self._params = params or {}
self._collected_targets: Dict[int, List["StageInfo"]] = {}
self._metrics: Dict[str, List[str]] = {}
self._plots: Dict[str, Any] = {}
self._params: Dict[str, List[str]] = {}
@cached_property
def stages(self) -> List["Stage"]: # pylint: disable=method-hidden
# note that ideally we should be keeping this in a set as it is unique,
# hashable and has no concept of orderliness on its own. But we depend
# on this to be somewhat ordered for status/metrics/plots, etc.
return self._collect()
@cached_property
def _top_metrics(self):
self._collect()
return self._metrics
@cached_property
def _top_plots(self):
self._collect()
return self._plots
@cached_property
def _plot_sources(self):
from dvc.repo.plots import _collect_pipeline_files
sources = []
for data in _collect_pipeline_files(self.repo, [], {}).values():
for plot_id, props in data.get("data", {}).items():
if isinstance(props.get("y"), dict):
sources.extend(props["y"])
if isinstance(props.get("x"), dict):
sources.extend(props["x"])
else:
sources.append(plot_id)
return sources
@cached_property
def _top_params(self):
self._collect()
return self._params
def _collect(self):
if "stages" in self.__dict__:
return self.stages
onerror = self.repo.stage_collection_error_handler
# pylint: disable=protected-access
(
stages,
metrics,
plots,
params,
) = self.stage_collector._collect_all_from_repo(onerror=onerror)
self.stages = stages
self._metrics = metrics
self._plots = plots
self._params = params
return stages
def __repr__(self) -> str:
from dvc.fs import LocalFileSystem
rev = "workspace"
if not isinstance(self.fs, LocalFileSystem):
if not isinstance(self.repo.fs, LocalFileSystem):
rev = self.repo.get_rev()[:7]
return f"Index({self.repo}, fs@{rev})"
@staticmethod
def _hash_targets(
targets: Iterable[Optional[str]],
**kwargs: Any,
) -> int:
return hash(
(
frozenset(targets),
kwargs.get("with_deps", False),
kwargs.get("recursive", False),
)
@classmethod
def from_repo(
cls, repo: "Repo", onerror: Callable[[str, Exception], None] = None
) -> "Index":
stages = []
metrics = {}
plots = {}
params = {}
onerror = onerror or repo.stage_collection_error_handler
for _, idx in collect_files(repo, onerror=onerror):
# pylint: disable=protected-access
stages.extend(idx.stages)
metrics.update(idx._metrics)
plots.update(idx._plots)
params.update(idx._params)
return cls(
repo,
stages=stages,
metrics=metrics,
plots=plots,
params=params,
)
def collect_targets(
self, targets: Optional["TargetType"], **kwargs: Any
) -> List["StageInfo"]:
from itertools import chain
@classmethod
def from_file(cls, repo: "Repo", path: str) -> "Index":
from dvc.dvcfile import make_dvcfile
from dvc.repo.stage import StageInfo
from dvc.utils.collections import ensure_list
dvcfile = make_dvcfile(repo, path)
return cls(
repo,
stages=list(dvcfile.stages.values()),
metrics={path: dvcfile.metrics} if dvcfile.metrics else {},
plots={path: dvcfile.plots} if dvcfile.plots else {},
params={path: dvcfile.params} if dvcfile.params else {},
)
targets = ensure_list(targets)
if not targets:
return [StageInfo(stage) for stage in self.stages]
targets_hash = self._hash_targets(targets, **kwargs)
if targets_hash not in self._collected_targets:
self._collected_targets[targets_hash] = list(
chain.from_iterable(
self.stage_collector.collect_granular(target, **kwargs)
for target in targets
)
)
return self._collected_targets[targets_hash]
def update(self, stages: Iterable["Stage"]) -> "Index":
stages = set(stages)
# we remove existing stages with same hashes at first
# and then re-add the new ones later.
stages_set = (set(self.stages) - stages) | stages
return self.__class__(
self.repo,
stages=list(stages_set),
metrics=self._metrics,
plots=self._plots,
params=self._params,
)
def targets_view(
self,
targets: Optional["TargetType"],
stage_filter: Optional[Callable[["Stage"], bool]] = None,
outs_filter: Optional[Callable[["Output"], bool]] = None,
**kwargs: Any,
) -> "IndexView":
"""Return read-only view of index for the specified targets.
@cached_property
def outs_trie(self) -> "Trie":
from dvc.repo.trie import build_outs_trie
Args:
targets: Targets to collect
stage_filter: Optional stage filter to be applied after collecting
targets.
outs_filter: Optional output filter to be applied after collecting
targets.
return build_outs_trie(self.stages)
Additional kwargs will be passed into the stage collector.
@cached_property
def outs_graph(self) -> "DiGraph":
from dvc.repo.graph import build_outs_graph
Note:
If both stage_filter and outs_filter are provided, stage_filter
will be applied first, and the resulting view will only contain
outputs from stages that matched stage_filter. Outputs from stages
that did not match will be excluded from the view (whether or not
the output would have matched outs_filter).
"""
stage_infos = [
stage_info
for stage_info in self.collect_targets(targets, **kwargs)
if not stage_filter or stage_filter(stage_info.stage)
]
return IndexView(self, stage_infos, outs_filter=outs_filter)
return build_outs_graph(self.graph, self.outs_trie)
@cached_property
def graph(self) -> "DiGraph":
from dvc.repo.graph import build_graph
return build_graph(self.stages, self.outs_trie)
def check_graph(self) -> None:
if not getattr(self.repo, "_skip_graph_checks", False):
self.graph # pylint: disable=pointless-statement
@property
def params(self) -> Iterator["ParamsDependency"]:
from dvc.dependency import ParamsDependency
for dep in self.deps:
if isinstance(dep, ParamsDependency):
yield dep
@property
def outs(self) -> Iterator["Output"]:
@@ -228,31 +239,20 @@ class Index:
for stage in self.stages:
yield from stage.deps
@property
def params(self) -> Iterator["ParamsDependency"]:
from dvc.dependency import ParamsDependency
for dep in self.deps:
if isinstance(dep, ParamsDependency):
yield dep
@cached_property
def outs_trie(self) -> "Trie":
from dvc.repo.trie import build_outs_trie
def _plot_sources(self) -> List[str]:
from dvc.repo.plots import _collect_pipeline_files
return build_outs_trie(self.stages)
@cached_property
def graph(self) -> "DiGraph":
from dvc.repo.graph import build_graph
return build_graph(self.stages, self.outs_trie)
@cached_property
def outs_graph(self) -> "DiGraph":
from dvc.repo.graph import build_outs_graph
return build_outs_graph(self.graph, self.outs_trie)
sources: List[str] = []
for data in _collect_pipeline_files(self.repo, [], {}).values():
for plot_id, props in data.get("data", {}).items():
if isinstance(props.get("y"), dict):
sources.extend(props["y"])
if isinstance(props.get("x"), dict):
sources.extend(props["x"])
else:
sources.append(plot_id)
return sources
@cached_property
def data(self) -> "Dict[str, DataIndex]":
@@ -278,6 +278,40 @@ class Index:
return dict(by_workspace)
@staticmethod
def _hash_targets(
targets: Iterable[Optional[str]],
**kwargs: Any,
) -> int:
return hash(
(
frozenset(targets),
kwargs.get("with_deps", False),
kwargs.get("recursive", False),
)
)
def collect_targets(
self, targets: Optional["TargetType"], **kwargs: Any
) -> List["StageInfo"]:
from itertools import chain
from dvc.repo.stage import StageInfo
from dvc.utils.collections import ensure_list
targets = ensure_list(targets)
if not targets:
return [StageInfo(stage) for stage in self.stages]
targets_hash = self._hash_targets(targets, **kwargs)
if targets_hash not in self._collected_targets:
self._collected_targets[targets_hash] = list(
chain.from_iterable(
self.repo.stage.collect_granular(target, **kwargs)
for target in targets
)
)
return self._collected_targets[targets_hash]
def used_objs(
self,
targets: "TargetType" = None,
@@ -305,23 +339,34 @@ class Index:
used[odb].update(objs)
return used
# Following methods help us treat the collection as a set-like structure
# and provides faux-immutability.
# These methods do not preserve stages order.
def update(self, stages: Iterable["Stage"]) -> "Index":
new_stages = set(stages)
# we remove existing stages with same hashes at first
# and then re-add the new ones later.
stages_set = (set(self.stages) - new_stages) | new_stages
return Index(self.repo, self.fs, stages=list(stages_set))
def add(self, stage: "Stage") -> "Index":
return self.update([stage])
def check_graph(self) -> None:
if not getattr(self.repo, "_skip_graph_checks", False):
self.graph # pylint: disable=pointless-statement
def targets_view(
self,
targets: Optional["TargetType"],
stage_filter: Optional[Callable[["Stage"], bool]] = None,
outs_filter: Optional[Callable[["Output"], bool]] = None,
**kwargs: Any,
) -> "IndexView":
"""Return read-only view of index for the specified targets.
Args:
targets: Targets to collect
stage_filter: Optional stage filter to be applied after collecting
targets.
outs_filter: Optional output filter to be applied after collecting
targets.
Additional kwargs will be passed into the stage collector.
Note:
If both stage_filter and outs_filter are provided, stage_filter
will be applied first, and the resulting view will only contain
outputs from stages that matched stage_filter. Outputs from stages
that did not match will be excluded from the view (whether or not
the output would have matched outs_filter).
"""
stage_infos = [
stage_info
for stage_info in self.collect_targets(targets, **kwargs)
if not stage_filter or stage_filter(stage_info.stage)
]
return IndexView(self, stage_infos, outs_filter=outs_filter)
class _DataPrefixes(NamedTuple):
@@ -342,13 +387,9 @@ class IndexView:
self._stage_infos = stage_infos
# NOTE: stage_infos might have the same stage multiple times but with
# different filter_info
self._stages = list({stage for stage, _ in stage_infos})
self.stages = list({stage for stage, _ in stage_infos})
self._outs_filter = outs_filter
@property
def stages(self):
return self._stages
@property
def deps(self) -> Iterator["Dependency"]:
for stage in self.stages:
@@ -385,8 +426,6 @@ class IndexView:
@cached_property
def data(self) -> Dict[str, Union["DataIndex", "DataIndexView"]]:
from functools import partial
from dvc_data.index import DataIndex, view
def key_filter(workspace: str, key: "DataIndexKey"):
@@ -401,7 +440,7 @@ class IndexView:
data: Dict[str, Union["DataIndex", "DataIndexView"]] = {}
for workspace, data_index in self._index.data.items():
if self._stages:
if self.stages:
data[workspace] = view(
data_index, partial(key_filter, workspace)
)
dvc/repo/metrics/show.py
+1 -1
View File
@@ -34,7 +34,7 @@ def _to_fs_paths(metrics: List[Output]) -> StrPaths:
def _collect_top_level_metrics(repo):
top_metrics = repo.index._top_metrics # pylint: disable=protected-access
top_metrics = repo.index._metrics # pylint: disable=protected-access
for dvcfile, metrics in top_metrics.items():
wdir = repo.fs.path.relpath(
repo.fs.path.parent(dvcfile), repo.root_dir
dvc/repo/params/show.py
+1 -1
View File
@@ -41,7 +41,7 @@ def _is_params(dep: "Output"):
def _collect_top_level_params(repo):
top_params = repo.index._top_params # pylint: disable=protected-access
top_params = repo.index._params # pylint: disable=protected-access
for dvcfile, params in top_params.items():
wdir = repo.fs.path.relpath(
repo.fs.path.parent(dvcfile), repo.root_dir
dvc/repo/plots/__init__.py
+1 -1
View File
@@ -453,7 +453,7 @@ def _resolve_definitions(
def _collect_pipeline_files(repo, targets: List[str], props, onerror=None):
result: Dict[str, Dict] = {}
top_plots = repo.index._top_plots # pylint: disable=protected-access
top_plots = repo.index._plots # pylint: disable=protected-access
for dvcfile, plots_def in top_plots.items():
dvcfile_path = _relpath(repo.dvcfs, dvcfile)
dvcfile_defs_dict: Dict[str, Union[Dict, None]] = {}
dvc/repo/stage.py
+12 -195
View File
@@ -1,30 +1,14 @@
import fnmatch
import logging
import time
import typing
from contextlib import suppress
from functools import partial, wraps
from typing import (
Callable,
Iterable,
List,
NamedTuple,
Optional,
Set,
Tuple,
Union,
)
from functools import wraps
from typing import Iterable, List, NamedTuple, Optional, Set, Tuple, Union
from funcy.debug import format_time
from dvc.exceptions import (
DvcException,
NoOutputOrStageError,
OutputNotFoundError,
)
from dvc.exceptions import NoOutputOrStageError, OutputNotFoundError
from dvc.repo import lock_repo
from dvc.ui import ui
from dvc.utils import as_posix, parse_target, relpath
from dvc.utils import as_posix, parse_target
logger = logging.getLogger(__name__)
@@ -39,14 +23,6 @@ if typing.TYPE_CHECKING:
PIPELINE_FILE = "dvc.yaml"
def log_walk(seq):
for root, dirs, files in seq:
start = time.perf_counter()
yield root, dirs, files
duration = format_time(time.perf_counter() - start)
logger.trace("%s in collecting stages from %s", duration, root)
class StageInfo(NamedTuple):
stage: "Stage"
filter_info: Optional[str] = None
@@ -75,7 +51,9 @@ def _maybe_collect_from_dvc_yaml(
if loader.fs.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
if with_deps:
return _collect_with_deps(stages, loader.repo.index.graph)
return stages
def _collect_specific_target(
@@ -117,9 +95,9 @@ def locked(f):
class StageLoad:
def __init__(self, repo: "Repo", fs=None) -> None:
def __init__(self, repo: "Repo") -> None:
self.repo: "Repo" = repo
self._fs = fs
self.fs = repo.fs
@locked
def add(
@@ -207,8 +185,7 @@ class StageLoad:
check_stage_exists(self.repo, stage, stage.path)
new_index = self.repo.index.add(stage)
new_index.check_graph()
self.repo.ensure_graph_correctness_with(stages={stage})
restore_fields(stage)
return stage
@@ -323,16 +300,6 @@ class StageLoad:
"""Load stages from `path`, filtered with `expr` provided."""
return self.load_all(path, expr, glob=True)
@property
def fs(self):
if self._fs:
return self._fs
return self.repo.fs
@property
def graph(self) -> "DiGraph":
return self.repo.index.graph
def collect(
self,
target: str = None,
@@ -375,13 +342,13 @@ class StageLoad:
from dvc.repo.graph import collect_inside_path
path = self.fs.path.abspath(target)
return collect_inside_path(path, graph or self.graph)
return collect_inside_path(path, graph or self.repo.index.graph)
stages = self.from_target(target, glob=glob)
if not with_deps:
return stages
return _collect_with_deps(stages, graph or self.graph)
return _collect_with_deps(stages, graph or self.repo.index.graph)
def collect_granular(
self,
@@ -449,153 +416,3 @@ class StageLoad:
raise NoOutputOrStageError(target, exc.file) from exc
return [StageInfo(stage) for stage in stages]
def _collect_repo(self, onerror: Callable[[str, Exception], None] = None):
"""Collects all of the stages present in the DVC repo.
Args:
onerror (optional): callable that will be called with two args:
the filepath whose collection failed and the exc instance.
It can report the error to continue with the collection
(and, skip failed ones), or raise the exception to abort
the collection.
"""
from dvc.dvcfile import is_valid_filename
from dvc.fs import LocalFileSystem
scm = self.repo.scm
sep = self.fs.sep
outs: Set[str] = set()
is_local_fs = isinstance(self.fs, LocalFileSystem)
def is_ignored(path):
# apply only for the local fs
return is_local_fs and scm.is_ignored(path)
def is_dvcfile_and_not_ignored(root, file):
return is_valid_filename(file) and not is_ignored(
f"{root}{sep}{file}"
)
def is_out_or_ignored(root, directory):
dir_path = f"{root}{sep}{directory}"
# trailing slash needed to check if a directory is gitignored
return dir_path in outs or is_ignored(f"{dir_path}{sep}")
walk_iter = self.repo.dvcignore.walk(self.fs, self.repo.root_dir)
if logger.isEnabledFor(logging.TRACE): # type: ignore[attr-defined]
walk_iter = log_walk(walk_iter)
for root, dirs, files in walk_iter:
dvcfile_filter = partial(is_dvcfile_and_not_ignored, root)
for file in filter(dvcfile_filter, files):
file_path = self.fs.path.join(root, file)
try:
new_stages = self.load_file(file_path)
except DvcException as exc:
if onerror:
onerror(relpath(file_path), exc)
continue
raise
yield from new_stages
outs.update(
out.fspath
for stage in new_stages
for out in stage.outs
if out.protocol == "local"
)
dirs[:] = [d for d in dirs if not is_out_or_ignored(root, d)]
def collect_repo(self, onerror: Callable[[str, Exception], None] = None):
return list(self._collect_repo(onerror))
def _load_file(self, path):
from dvc.dvcfile import Dvcfile
from dvc.stage.loader import SingleStageLoader, StageLoader
path = self._get_filepath(path)
dvcfile = Dvcfile(self.repo, path)
stages = dvcfile.stages
if isinstance(stages, SingleStageLoader):
stages_ = [stages[None]]
else:
assert isinstance(stages, StageLoader)
keys = self._get_keys(stages)
stages_ = [stages[key] for key in keys]
return stages_, dvcfile
def _collect_all_from_repo( # noqa: C901
self, onerror: Callable[[str, Exception], None] = None
):
"""Collects all of the stages present in the DVC repo.
Args:
onerror (optional): callable that will be called with two args:
the filepath whose collection failed and the exc instance.
It can report the error to continue with the collection
(and, skip failed ones), or raise the exception to abort
the collection.
"""
from dvc.dvcfile import is_valid_filename
from dvc.fs import LocalFileSystem
scm = self.repo.scm
sep = self.fs.sep
outs: Set[str] = set()
is_local_fs = isinstance(self.fs, LocalFileSystem)
def is_ignored(path):
# apply only for the local fs
return is_local_fs and scm.is_ignored(path)
def is_dvcfile_and_not_ignored(root, file):
return is_valid_filename(file) and not is_ignored(
f"{root}{sep}{file}"
)
def is_out_or_ignored(root, directory):
dir_path = f"{root}{sep}{directory}"
# trailing slash needed to check if a directory is gitignored
return dir_path in outs or is_ignored(f"{dir_path}{sep}")
walk_iter = self.repo.dvcignore.walk(self.fs, self.repo.root_dir)
if logger.isEnabledFor(logging.TRACE): # type: ignore[attr-defined]
walk_iter = log_walk(walk_iter)
stages = []
metrics = {}
plots = {}
params = {}
for root, dirs, files in walk_iter:
dvcfile_filter = partial(is_dvcfile_and_not_ignored, root)
for file in filter(dvcfile_filter, files):
file_path = self.fs.path.join(root, file)
try:
new_stages, dvcfile = self._load_file(file_path)
except DvcException as exc:
if onerror:
onerror(relpath(file_path), exc)
continue
raise
stages.extend(new_stages)
if new_metrics := dvcfile.metrics:
metrics[file_path] = new_metrics
if new_plots := dvcfile.plots:
plots[file_path] = new_plots
if new_params := dvcfile.params:
params[file_path] = new_params
outs.update(
out.fspath
for stage in new_stages
for out in stage.outs
if out.protocol == "local"
)
dirs[:] = [d for d in dirs if not is_out_or_ignored(root, d)]
return stages, metrics, plots, params
tests/func/test_remove.py
+1 -1
View File
@@ -26,7 +26,7 @@ def test_remove(tmp_dir, scm, dvc, run_copy, remove_outs):
for stage in [stage1, stage2, stage3]:
dvc.remove(stage.addressing, outs=remove_outs)
out_exists = (out.exists for out in stage.outs)
assert stage not in dvc.stage.collect_repo()
assert stage not in dvc.index.stages
if remove_outs:
assert not any(out_exists)
else:
tests/func/test_repo_index.py
+33 -14
View File
@@ -12,11 +12,9 @@ def test_index(tmp_dir, scm, dvc, run_copy):
stage2 = run_copy("foo", "bar", name="copy-foo-bar")
tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add")
index = Index(dvc)
assert index.fs == dvc.fs
index = Index.from_repo(dvc)
assert set(index.stages) == {stage1, stage2}
assert index.outs_graph
assert index.graph
assert isinstance(index.outs_trie, Trie)
@@ -70,7 +68,7 @@ def test_deps_outs_getters(tmp_dir, dvc, run_copy_metrics):
name="copy-metrics2",
)
index = Index(dvc)
index = Index.from_repo(dvc)
stages = [foo_stage, run_stage1, run_stage2]
(metrics,) = run_stage1.outs
@@ -88,16 +86,16 @@ def test_deps_outs_getters(tmp_dir, dvc, run_copy_metrics):
assert outputs_equal(index.params, [params])
def test_add_update(dvc):
"""Test that add/update overwrites existing stages with the new ones.
def test_update(dvc):
"""Test that update overwrites existing stages with the new ones.
The old stages and the new ones might have same hash, so we are
making sure that the old stages were removed and replaced by new ones
using `id`/`is` checks.
"""
index = Index(dvc)
index = Index.from_repo(dvc)
new_stage = Stage(dvc, path="path1")
new_index = index.add(new_stage)
new_index = index.update({new_stage})
assert not index.stages
assert new_index.stages == [new_stage]
@@ -151,14 +149,14 @@ def get_index(dvc, rev):
if rev != "workspace":
assert next(brancher) == "workspace"
next(brancher)
return Index(dvc)
return Index.from_repo(dvc)
@pytest.mark.parametrize("rev", ["workspace", "HEAD"])
def test_used_objs(tmp_dir, scm, dvc, run_copy, rev):
from dvc_data.hashfile.hash_info import HashInfo
dvc.config["core"]["autostage"] = True
dvc.scm_context.autostage = True
tmp_dir.dvc_gen({"dir": {"subdir": {"file": "file"}}, "foo": "foo"})
run_copy("foo", "bar", name="copy-foo-bar")
scm.commit("commit")
@@ -196,7 +194,7 @@ def test_view_granular_dir(tmp_dir, scm, dvc, run_copy):
{"dir": {"subdir": {"in-subdir": "in-subdir"}, "in-dir": "in-dir"}},
commit="init",
)
index = Index(dvc)
index = Index.from_repo(dvc)
# view should include the specific target, parent dirs, and children
# view should exclude any siblings of the target
@@ -218,7 +216,7 @@ def test_view_stage_filter(tmp_dir, scm, dvc, run_copy):
(stage1,) = tmp_dir.dvc_gen("foo", "foo")
stage2 = run_copy("foo", "bar", name="copy-foo-bar")
tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add")
index = Index(dvc)
index = Index.from_repo(dvc)
view = index.targets_view(None)
assert set(view.stages) == {stage1, stage2}
@@ -239,7 +237,7 @@ def test_view_outs_filter(tmp_dir, scm, dvc, run_copy):
(stage1,) = tmp_dir.dvc_gen("foo", "foo")
stage2 = run_copy("foo", "bar", name="copy-foo-bar")
tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add")
index = Index(dvc)
index = Index.from_repo(dvc)
view = index.targets_view(None, outs_filter=lambda o: o.def_path == "foo")
assert set(view.stages) == {stage1, stage2}
@@ -252,7 +250,7 @@ def test_view_combined_filter(tmp_dir, scm, dvc, run_copy):
(stage1,) = tmp_dir.dvc_gen("foo", "foo")
stage2 = run_copy("foo", "bar", name="copy-foo-bar")
tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add")
index = Index(dvc)
index = Index.from_repo(dvc)
view = index.targets_view(
None,
@@ -271,3 +269,24 @@ def test_view_combined_filter(tmp_dir, scm, dvc, run_copy):
assert {out.fs_path for out in view.outs} == {
out.fs_path for out in stage2.outs
}
def test_with_gitignore(tmp_dir, dvc, scm):
(stage,) = tmp_dir.dvc_gen({"data": {"foo": "foo", "bar": "bar"}})
index = Index.from_repo(dvc)
assert index.stages == [stage]
scm.ignore(stage.path)
scm._reset()
index = Index.from_repo(dvc)
assert not index.stages
def test_ignored_dir_unignored_pattern(tmp_dir, dvc, scm):
tmp_dir.gen({".gitignore": "data/**\n!data/**/\n!data/**/*.dvc"})
scm.add([".gitignore"])
(stage,) = tmp_dir.dvc_gen({"data/raw/tracked.csv": "5,6,7,8"})
index = Index.from_repo(dvc)
assert index.stages == [stage]
tests/func/test_stage.py
-7
View File
@@ -239,13 +239,6 @@ def test_parent_repo_collect_stages(tmp_dir, scm, dvc):
assert deep_subrepo_stages != []
def test_collect_repo_ignored_dir_unignored_pattern(tmp_dir, dvc, scm):
tmp_dir.gen({".gitignore": "data/**\n!data/**/\n!data/**/*.dvc"})
scm.add([".gitignore"])
(stage,) = tmp_dir.dvc_gen({"data/raw/tracked.csv": "5,6,7,8"})
assert dvc.stage.collect_repo() == [stage]
@pytest.mark.parametrize("with_deps", (False, True))
def test_collect_symlink(tmp_dir, dvc, with_deps):
tmp_dir.gen({"data": {"foo": "foo contents"}})
tests/func/test_stage_load.py
+2 -13
View File
@@ -391,7 +391,7 @@ def test_collect_optimization(tmp_dir, dvc, mocker):
# Forget cached stages and graph and error out on collection
dvc._reset()
mocker.patch(
"dvc.repo.index.Index.stages",
"dvc.repo.Repo.index",
property(raiser(Exception("Should not collect"))),
)
@@ -406,7 +406,7 @@ def test_collect_optimization_on_stage_name(tmp_dir, dvc, mocker, run_copy):
# Forget cached stages and graph and error out on collection
dvc._reset()
mocker.patch(
"dvc.repo.index.Index.stages",
"dvc.repo.Repo.index",
property(raiser(Exception("Should not collect"))),
)
@@ -431,17 +431,6 @@ def test_collect_repo_callback(tmp_dir, dvc, mocker):
assert isinstance(exc, YAMLValidationError)
def test_gitignored_collect_repo(tmp_dir, dvc, scm):
(stage,) = tmp_dir.dvc_gen({"data": {"foo": "foo", "bar": "bar"}})
assert dvc.stage.collect_repo() == [stage]
scm.ignore(stage.path)
scm._reset()
assert not dvc.stage.collect_repo()
def test_gitignored_file_try_collect_granular_for_data_files(
tmp_dir, dvc, scm
):
tests/unit/command/test_status.py
+7 -7
View File
@@ -76,12 +76,14 @@ def test_status_quiet(dvc, mocker, caplog, capsys, status, ret):
def test_status_empty(dvc, mocker, capsys):
from dvc.repo.index import Index
cli_args = parse_args(["status"])
assert cli_args.func == CmdDataStatus
cmd = cli_args.func(cli_args)
spy = mocker.spy(cmd.repo.stage, "_collect_all_from_repo")
spy = mocker.spy(Index, "from_repo")
assert cmd.run() == 0
@@ -100,6 +102,8 @@ def test_status_empty(dvc, mocker, capsys):
],
)
def test_status_up_to_date(dvc, mocker, capsys, cloud_opts, expected_message):
from dvc.repo.index import Index
cli_args = parse_args(["status", *cloud_opts])
assert cli_args.func == CmdDataStatus
@@ -107,12 +111,8 @@ def test_status_up_to_date(dvc, mocker, capsys, cloud_opts, expected_message):
mocker.patch.dict(cmd.repo.config, {"core": {"remote": "default"}})
mocker.patch.object(cmd.repo, "status", autospec=True, return_value={})
mocker.patch.object(
cmd.repo.stage,
"_collect_all_from_repo",
return_value=[[object()], {}, {}, {}],
autospec=True,
)
mocker.patch("dvc.repo.Repo.index", return_value=Index(dvc, [object()]))
cmd.repo._reset = mocker.Mock()
assert cmd.run() == 0
captured = capsys.readouterr()