diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 81d4ef8d6..9fb3f5ee1 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -179,28 +179,34 @@ class Repo(object): import networkx as nx from dvc.stage import Stage - G = graph or self.graph - if not target: - return list(G) + return list(graph) if graph else self.stages target = os.path.abspath(target) if recursive and os.path.isdir(target): - stages = nx.dfs_postorder_nodes(G) + stages = nx.dfs_postorder_nodes(graph or self.graph) return [stage for stage in stages if path_isin(stage.path, target)] stage = Stage.load(self, target) + + # Optimization: do not collect the graph for a specific target if not with_deps: return [stage] - pipeline = get_pipeline(get_pipelines(G), stage) + pipeline = get_pipeline(get_pipelines(graph or self.graph), stage) return list(nx.dfs_postorder_nodes(pipeline, stage)) def collect_granular(self, target, *args, **kwargs): + from dvc.stage import Stage + if not target: return [(stage, None) for stage in self.stages] + # Optimization: do not collect the graph for a specific .dvc target + if Stage.is_valid_filename(target) and not kwargs.get("with_deps"): + return [(Stage.load(self, target), None)] + try: (out,) = self.find_outs_by_path(target, strict=False) filter_info = PathInfo(os.path.abspath(target)) diff --git a/tests/unit/repo/test_repo.py b/tests/unit/repo/test_repo.py index b6ee260de..4638ff011 100644 --- a/tests/unit/repo/test_repo.py +++ b/tests/unit/repo/test_repo.py @@ -1,5 +1,6 @@ import os +from funcy import raiser import pytest from dvc.repo import locked @@ -66,3 +67,18 @@ def test_locked(mocker): mocker.call.method(repo, args, kwargs), mocker.call._reset(), ] + + +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)