diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py index e7d5055e0..88ebee36d 100644 --- a/dvc/dvcfile.py +++ b/dvc/dvcfile.py @@ -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) diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 93452910f..bd8edfb11 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -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): diff --git a/dvc/repo/add.py b/dvc/repo/add.py index 5262923fb..be53fac56 100644 --- a/dvc/repo/add.py +++ b/dvc/repo/add.py @@ -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: diff --git a/dvc/repo/imp_url.py b/dvc/repo/imp_url.py index eeb0abb18..144a84d76 100644 --- a/dvc/repo/imp_url.py +++ b/dvc/repo/imp_url.py @@ -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} diff --git a/dvc/repo/index.py b/dvc/repo/index.py index c988b1f0b..2b45ce4ac 100644 --- a/dvc/repo/index.py +++ b/dvc/repo/index.py @@ -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) ) diff --git a/dvc/repo/metrics/show.py b/dvc/repo/metrics/show.py index a7b4a1158..c1b7fb73b 100644 --- a/dvc/repo/metrics/show.py +++ b/dvc/repo/metrics/show.py @@ -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 diff --git a/dvc/repo/params/show.py b/dvc/repo/params/show.py index 58d13e31e..22669fbee 100644 --- a/dvc/repo/params/show.py +++ b/dvc/repo/params/show.py @@ -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 diff --git a/dvc/repo/plots/__init__.py b/dvc/repo/plots/__init__.py index 0e81b0258..43ed2b07c 100644 --- a/dvc/repo/plots/__init__.py +++ b/dvc/repo/plots/__init__.py @@ -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]] = {} diff --git a/dvc/repo/stage.py b/dvc/repo/stage.py index ec6702994..48e221e0b 100644 --- a/dvc/repo/stage.py +++ b/dvc/repo/stage.py @@ -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 diff --git a/tests/func/test_remove.py b/tests/func/test_remove.py index 25313ac7d..40ad64110 100644 --- a/tests/func/test_remove.py +++ b/tests/func/test_remove.py @@ -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: diff --git a/tests/func/test_repo_index.py b/tests/func/test_repo_index.py index f5e965fd3..6553ddf8f 100644 --- a/tests/func/test_repo_index.py +++ b/tests/func/test_repo_index.py @@ -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] diff --git a/tests/func/test_stage.py b/tests/func/test_stage.py index 6596b15d5..657c028dc 100644 --- a/tests/func/test_stage.py +++ b/tests/func/test_stage.py @@ -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"}}) diff --git a/tests/func/test_stage_load.py b/tests/func/test_stage_load.py index cedae6b87..49c78c0da 100644 --- a/tests/func/test_stage_load.py +++ b/tests/func/test_stage_load.py @@ -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 ): diff --git a/tests/unit/command/test_status.py b/tests/unit/command/test_status.py index fd69a84f5..9634e0081 100644 --- a/tests/unit/command/test_status.py +++ b/tests/unit/command/test_status.py @@ -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()