diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 41daa23a7..19b1e4f50 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -8,7 +8,6 @@ from funcy import cached_property, cat from git import InvalidGitRepositoryError from dvc.config import Config -from dvc.dvcfile import is_valid_filename from dvc.exceptions import FileMissingError from dvc.exceptions import IsADirectoryError as DvcIsADirectoryError from dvc.exceptions import NotDvcRepoError, OutputNotFoundError @@ -183,6 +182,7 @@ class Repo: self.metrics = Metrics(self) self.plots = Plots(self) self.params = Params(self) + self.stage_collection_error_handler = None self._lock_depth = 0 @cached_property @@ -371,25 +371,8 @@ class Repo: NOTE: For large repos, this could be an expensive operation. Consider using some memoization. """ - return self._collect_stages() - - def _collect_stages(self): - stages = [] - outs = set() - - for root, dirs, files in self.tree.walk(self.root_dir): - for file_name in filter(is_valid_filename, files): - file_path = os.path.join(root, file_name) - new_stages = self.stage.load_file(file_path) - stages.extend(new_stages) - outs.update( - out.fspath - for stage in new_stages - for out in stage.outs - if out.scheme == "local" - ) - dirs[:] = [d for d in dirs if os.path.join(root, d) not in outs] - return stages + error_handler = self.stage_collection_error_handler + return self.stage.collect_repo(onerror=error_handler) def find_outs_by_path(self, path, outs=None, recursive=False, strict=True): if not outs: diff --git a/dvc/repo/stage.py b/dvc/repo/stage.py index a5043eb04..eaf137b0c 100644 --- a/dvc/repo/stage.py +++ b/dvc/repo/stage.py @@ -6,11 +6,15 @@ from contextlib import suppress from typing import Iterable, List, NamedTuple, Optional, Set, Tuple from dvc.dvcfile import PIPELINE_FILE, Dvcfile, is_valid_filename -from dvc.exceptions import NoOutputOrStageError, OutputNotFoundError +from dvc.exceptions import ( + DvcException, + 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 +from dvc.utils import parse_target, relpath logger = logging.getLogger(__name__) @@ -316,3 +320,37 @@ class StageLoad: raise NoOutputOrStageError(target, exc.file) from exc return [StageInfo(stage) for stage in stages] + + def collect_repo(self, onerror=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. + The decision to raise error is with the callback. + Otherwise, the file will be skipped, and dvc will partially + load the stages. + """ + stages = [] + outs = set() + for root, dirs, files in self.tree.walk(self.repo.root_dir): + for file_name in filter(is_valid_filename, files): + file_path = os.path.join(root, file_name) + + try: + new_stages = self.load_file(file_path) + except DvcException as exc: + if onerror: + onerror(relpath(file_path), exc) + continue + raise + + stages.extend(new_stages) + outs.update( + out.fspath + for stage in new_stages + for out in stage.outs + if out.scheme == "local" + ) + dirs[:] = [d for d in dirs if os.path.join(root, d) not in outs] + return stages diff --git a/tests/func/test_remove.py b/tests/func/test_remove.py index e5f3ef4ff..0faeb7dbe 100644 --- a/tests/func/test_remove.py +++ b/tests/func/test_remove.py @@ -23,7 +23,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._collect_stages() + assert stage not in dvc.stage.collect_repo() if remove_outs: assert not any(out_exists) else: diff --git a/tests/func/test_stage_load.py b/tests/func/test_stage_load.py index b8d83608c..8f8e20c92 100644 --- a/tests/func/test_stage_load.py +++ b/tests/func/test_stage_load.py @@ -10,6 +10,7 @@ from dvc.path_info import PathInfo from dvc.repo import Repo from dvc.stage.exceptions import ( StageFileDoesNotExistError, + StageFileFormatError, StageNameUnspecified, StageNotFound, ) @@ -434,3 +435,19 @@ def test_collect_optimization_on_stage_name(tmp_dir, dvc, mocker, run_copy): # 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)] + + +def test_collect_repo_callback(tmp_dir, dvc, mocker): + mock = mocker.Mock() + dvc.stage_collection_error_handler = mock + + (stage,) = tmp_dir.dvc_gen("foo", "foo") + dump_yaml(tmp_dir / PIPELINE_FILE, {"stages": {"cmd": "echo hello world"}}) + + dvc._reset() + assert dvc.stages == [stage] + mock.assert_called_once() + + file_path, exc = mock.call_args[0] + assert file_path == PIPELINE_FILE + assert isinstance(exc, StageFileFormatError)