dvc: make graph consist of stages not their paths

dvc/command/pipeline.py

@@ -2,14 +2,11 @@ from __future__ import unicode_literals
 
 import argparse
 import logging
-
-from dvc.command.base import append_doc_link
-from dvc.command.base import CmdBase
-from dvc.command.base import fix_subparsers
-from dvc.exceptions import DvcException
-from dvc.utils import relpath
 from dvc.utils.compat import str
 
+from dvc.command.base import CmdBase, append_doc_link, fix_subparsers
+from dvc.exceptions import DvcException
+
 
 logger = logging.getLogger(__name__)
 
@@ -21,23 +18,21 @@ class CmdPipelineShow(CmdBase):
 
         stage = Stage.load(self.repo, target)
         G = self.repo.graph
-        stages = networkx.get_node_attributes(G, "stage")
-        node = relpath(stage.path, self.repo.root_dir)
-        nodes = networkx.dfs_postorder_nodes(G, node)
+        stages = networkx.dfs_postorder_nodes(G, stage)
 
         if locked:
-            nodes = [n for n in nodes if stages[n].locked]
+            stages = [s for s in stages if s.locked]
 
-        for n in nodes:
+        for stage in stages:
             if commands:
-                if stages[n].cmd is None:
+                if stage.cmd is None:
                     continue
-                logger.info(stages[n].cmd)
+                logger.info(stage.cmd)
             elif outs:
-                for out in stages[n].outs:
+                for out in stage.outs:
                     logger.info(str(out))
             else:
-                logger.info(n)
+                logger.info(stage.path_in_repo)
 
     def __build_graph(self, target, commands, outs):
         import networkx
@@ -45,14 +40,10 @@ class CmdPipelineShow(CmdBase):
         from dvc.repo.graph import get_pipeline
 
         stage = Stage.load(self.repo, target)
-        node = relpath(stage.path, self.repo.root_dir)
-
-        G = get_pipeline(self.repo.pipelines, node)
-        stages = networkx.get_node_attributes(G, "stage")
+        G = get_pipeline(self.repo.pipelines, stage)
 
         nodes = []
-        for n in G:
-            stage = stages[n]
+        for stage in G:
             if commands:
                 if stage.cmd is None:
                     continue
@@ -64,9 +55,7 @@ class CmdPipelineShow(CmdBase):
                 nodes.append(stage.relpath)
 
         edges = []
-        for e in G.edges():
-            from_stage = stages[e[0]]
-            to_stage = stages[e[1]]
+        for from_stage, to_stage in G.edges():
             if commands:
                 if to_stage.cmd is None:
                     continue
@@ -163,14 +152,11 @@ class CmdPipelineShow(CmdBase):
 
 class CmdPipelineList(CmdBase):
     def run(self):
-        import networkx
-
         pipelines = self.repo.pipelines
-        for p in pipelines:
-            stages = networkx.get_node_attributes(p, "stage")
-            for stage in stages:
-                logger.info(stage)
-            if len(stages) != 0:
+        for pipeline in pipelines:
+            for stage in pipeline:
+                logger.info(stage.relpath)
+            if len(pipeline) != 0:
                 logger.info("=" * 80)
         logger.info("{} pipelines total".format(len(pipelines)))
 

dvc/exceptions.py

@@ -39,17 +39,13 @@ class OutputDuplicationError(DvcException):
     """
 
     def __init__(self, output, stages):
-        assert isinstance(output, str) or isinstance(output, builtin_str)
-        assert isinstance(stages, list)
-        assert all(
-            isinstance(stage, str) or isinstance(stage, builtin_str)
-            for stage in stages
-        )
+        assert isinstance(output, (str, builtin_str))
+        assert all(hasattr(stage, "relpath") for stage in stages)
         msg = (
             "file/directory '{}' is specified as an output in more than one "
             "stage: {}\n"
             "This is not allowed. Consider using a different output name."
-        ).format(output, "\n    ".join(stages))
+        ).format(output, "\n    ".join(s.relpath for s in stages))
         super(OutputDuplicationError, self).__init__(msg)
 
 
@@ -75,21 +71,17 @@ class StagePathAsOutputError(DvcException):
     an output of another stage.
 
     Args:
-        cwd (str): path to the directory.
-        fname (str): path to the DVC-file that has cwd specified as an
-            output.
+        stage (Stage): a stage that is in some other stages output
+        output (str): an output covering the stage above
     """
 
-    def __init__(self, wdir, fname):
-        assert isinstance(wdir, str) or isinstance(wdir, builtin_str)
-        assert isinstance(fname, str) or isinstance(fname, builtin_str)
-        msg = (
-            "current working directory '{cwd}' is specified as an output in "
-            "'{fname}'. Use another CWD to prevent any data removal.".format(
-                cwd=wdir, fname=fname
+    def __init__(self, stage, output):
+        assert isinstance(output, (str, builtin_str))
+        super(StagePathAsOutputError, self).__init__(
+            "'{stage}' is within an output '{output}' of another stage".format(
+                stage=stage.relpath, output=output
             )
         )
-        super(StagePathAsOutputError, self).__init__(msg)
 
 
 class CircularDependencyError(DvcException):
@@ -167,7 +159,7 @@ class DvcParserError(DvcException):
 class CyclicGraphError(DvcException):
     def __init__(self, stages):
         assert isinstance(stages, list)
-        stages = "\n".join("\t- {}".format(stage) for stage in stages)
+        stages = "\n".join("\t- {}".format(stage.relpath) for stage in stages)
         msg = (
             "you've introduced a cycle in your pipeline that involves "
             "the following stages:"

dvc/repo/__init__.py

@@ -5,25 +5,22 @@ import os
 from contextlib import contextmanager
 from functools import wraps
 from itertools import chain
+from dvc.utils.compat import FileNotFoundError, fspath_py35, open as _open
 
 from funcy import cached_property
 
-from .graph import check_acyclic
-from .graph import get_pipeline
-from .graph import get_pipelines
-from .graph import get_stages
 from dvc.config import Config
-from dvc.exceptions import FileMissingError
-from dvc.exceptions import NotDvcRepoError
-from dvc.exceptions import OutputNotFoundError
+from dvc.exceptions import (
+    FileMissingError,
+    NotDvcRepoError,
+    OutputNotFoundError,
+)
 from dvc.ignore import DvcIgnoreFilter
 from dvc.path_info import PathInfo
 from dvc.remote.base import RemoteActionNotImplemented
 from dvc.utils import relpath
 from dvc.utils.fs import path_isin
-from dvc.utils.compat import FileNotFoundError
-from dvc.utils.compat import fspath_py35
-from dvc.utils.compat import open as _open
+from .graph import check_acyclic, get_pipeline, get_pipelines
 
 
 logger = logging.getLogger(__name__)
@@ -190,32 +187,20 @@ class Repo(object):
         G = graph or self.graph
 
         if not target:
-            return get_stages(G)
+            return list(G)
 
         target = os.path.abspath(target)
 
         if recursive and os.path.isdir(target):
-            attrs = nx.get_node_attributes(G, "stage")
-            nodes = [node for node in nx.dfs_postorder_nodes(G)]
-
-            ret = []
-            for node in nodes:
-                stage = attrs[node]
-                if path_isin(stage.path, target):
-                    ret.append(stage)
-            return ret
+            stages = nx.dfs_postorder_nodes(G)
+            return [stage for stage in stages if path_isin(stage.path, target)]
 
         stage = Stage.load(self, target)
         if not with_deps:
             return [stage]
 
-        node = relpath(stage.path, self.root_dir)
-        pipeline = get_pipeline(get_pipelines(G), node)
-
-        return [
-            pipeline.node[n]["stage"]
-            for n in nx.dfs_postorder_nodes(pipeline, node)
-        ]
+        pipeline = get_pipeline(get_pipelines(G), stage)
+        return list(nx.dfs_postorder_nodes(pipeline, stage))
 
     def used_cache(
         self,
@@ -323,15 +308,15 @@ class Repo(object):
         )
 
         G = nx.DiGraph()
-        stages = stages or self.collect_stages()
+        stages = stages or self.stages
         stages = [stage for stage in stages if stage]
         outs = {}
 
         for stage in stages:
             for out in stage.outs:
                 if out.path_info in outs:
-                    stages = [stage.relpath, outs[out.path_info].stage.relpath]
-                    raise OutputDuplicationError(str(out), stages)
+                    dup_stages = [stage, outs[out.path_info].stage]
+                    raise OutputDuplicationError(str(out), dup_stages)
                 outs[out.path_info] = out
 
         for stage in stages:
@@ -344,23 +329,19 @@ class Repo(object):
             stage_path_info = PathInfo(stage.path)
             for p in chain([stage_path_info], stage_path_info.parents):
                 if p in outs:
-                    raise StagePathAsOutputError(stage.wdir, stage.relpath)
+                    raise StagePathAsOutputError(stage, str(outs[p]))
 
         for stage in stages:
-            node = relpath(stage.path, self.root_dir)
-
-            G.add_node(node, stage=stage)
+            G.add_node(stage)
 
             for dep in stage.deps:
                 if dep.path_info is None:
                     continue
 
-                for out in outs:
-                    if out.overlaps(dep.path_info):
-                        dep_stage = outs[out].stage
-                        dep_node = relpath(dep_stage.path, self.root_dir)
-                        G.add_node(dep_node, stage=dep_stage)
-                        G.add_edge(node, dep_node)
+                for out_path_info, out in outs.items():
+                    if out_path_info.overlaps(dep.path_info):
+                        G.add_node(out.stage)
+                        G.add_edge(stage, out.stage)
 
         check_acyclic(G)
 
@@ -385,7 +366,8 @@ class Repo(object):
 
         return list(filter(filter_dirs, dirs))
 
-    def collect_stages(self):
+    @cached_property
+    def stages(self):
         """
         Walks down the root directory looking for Dvcfiles,
         skipping the directories that are related with
@@ -417,10 +399,6 @@ class Repo(object):
 
         return stages
 
-    @cached_property
-    def stages(self):
-        return get_stages(self.graph)
-
     def find_outs_by_path(self, path, outs=None, recursive=False):
         if not outs:
             outs = [out for stage in self.stages for out in stage.outs]

dvc/repo/graph.py

@@ -28,9 +28,3 @@ def get_pipelines(G):
     import networkx as nx
 
     return [G.subgraph(c).copy() for c in nx.weakly_connected_components(G)]
-
-
-def get_stages(G):
-    import networkx
-
-    return list(networkx.get_node_attributes(G, "stage").values())

dvc/repo/reproduce.py

@@ -2,20 +2,16 @@ from __future__ import unicode_literals
 
 import logging
 
-from . import locked
-from .graph import get_pipeline
-from .graph import get_pipelines
 from dvc.exceptions import ReproductionError
 from dvc.repo.scm_context import scm_context
-from dvc.utils import relpath
+from . import locked
+from .graph import get_pipeline, get_pipelines
 
 
 logger = logging.getLogger(__name__)
 
 
-def _reproduce_stage(stages, node, **kwargs):
-    stage = stages[node]
-
+def _reproduce_stage(stage, **kwargs):
     if stage.locked:
         logger.warning(
             "DVC-file '{path}' is locked. Its dependencies are"
@@ -36,16 +32,14 @@ def _get_active_graph(G):
     import networkx as nx
 
     active = G.copy()
-    stages = nx.get_node_attributes(G, "stage")
-    for node in G:
-        stage = stages[node]
+    for stage in G:
         if not stage.locked:
             continue
-        for n in nx.dfs_postorder_nodes(G, node):
-            if n == node:
+        for st in nx.dfs_postorder_nodes(G, stage):
+            if st == stage:
                 continue
-            if n in active:
-                active.remove_node(n)
+            if st in active:
+                active.remove_node(st)
     return active
 
 
@@ -59,7 +53,6 @@ def reproduce(
     all_pipelines=False,
     **kwargs
 ):
-    import networkx as nx
     from dvc.stage import Stage
 
     if not target and not all_pipelines:
@@ -81,39 +74,27 @@ def reproduce(
             pipelines = active_pipelines
         else:
             stage = Stage.load(self, target)
-            node = relpath(stage.path, self.root_dir)
-            pipelines = [get_pipeline(active_pipelines, node)]
+            pipelines = [get_pipeline(active_pipelines, stage)]
 
         targets = []
-        for G in pipelines:
-            attrs = nx.get_node_attributes(G, "stage")
-            for node in G:
-                if G.in_degree(node) == 0:
-                    targets.append(attrs[node])
+        for pipeline in pipelines:
+            for stage in pipeline:
+                if pipeline.in_degree(stage) == 0:
+                    targets.append(stage)
     else:
         targets = self.collect(target, recursive=recursive, graph=active_graph)
 
     ret = []
     for target in targets:
-        stages = _reproduce(self, active_graph, target, **kwargs)
+        stages = _reproduce_stages(active_graph, target, **kwargs)
         ret.extend(stages)
 
     return ret
 
 
-def _reproduce(self, G, stage, **kwargs):
-    import networkx as nx
-
-    stages = nx.get_node_attributes(G, "stage")
-    node = relpath(stage.path, self.root_dir)
-
-    return _reproduce_stages(G, stages, node, **kwargs)
-
-
 def _reproduce_stages(
     G,
-    stages,
-    node,
+    stage,
     downstream=False,
     ignore_build_cache=False,
     single_item=False,
@@ -153,11 +134,10 @@ def _reproduce_stages(
 
     The derived evaluation of _downstream_ B would be: [B, D, E]
     """
-
     import networkx as nx
 
     if single_item:
-        pipeline = [node]
+        pipeline = [stage]
     elif downstream:
         # NOTE (py3 only):
         # Python's `deepcopy` defaults to pickle/unpickle the object.
@@ -166,14 +146,14 @@ def _reproduce_stages(
         # We need to create a copy of the graph itself, and then reverse it,
         # instead of using graph.reverse() directly because it calls
         # `deepcopy` underneath -- unless copy=False is specified.
-        pipeline = nx.dfs_preorder_nodes(G.copy().reverse(copy=False), node)
+        pipeline = nx.dfs_preorder_nodes(G.copy().reverse(copy=False), stage)
     else:
-        pipeline = nx.dfs_postorder_nodes(G, node)
+        pipeline = nx.dfs_postorder_nodes(G, stage)
 
     result = []
-    for n in pipeline:
+    for st in pipeline:
         try:
-            ret = _reproduce_stage(stages, n, **kwargs)
+            ret = _reproduce_stage(st, **kwargs)
 
             if len(ret) != 0 and ignore_build_cache:
                 # NOTE: we are walking our pipeline from the top to the
@@ -183,7 +163,7 @@ def _reproduce_stages(
                 # dependencies didn't change.
                 kwargs["force"] = True
 
-            result += ret
+            result.extend(ret)
         except Exception as ex:
-            raise ReproductionError(stages[n].relpath, ex)
+            raise ReproductionError(st.relpath, ex)
     return result

dvc/stage.py

@@ -241,6 +241,20 @@ class Stage(object):
             path=self.relpath if self.path else "No path"
         )
 
+    def __hash__(self):
+        return hash(self.path_in_repo)
+
+    def __eq__(self, other):
+        return (
+            self.__class__ == other.__class__
+            and self.repo is other.repo
+            and self.path_in_repo == other.path_in_repo
+        )
+
+    @property
+    def path_in_repo(self):
+        return relpath(self.path, self.repo.root_dir)
+
     @property
     def relpath(self):
         return relpath(self.path)

tests/func/test_repo.py

@@ -32,7 +32,7 @@ def test_destroy(tmp_dir, dvc):
 def test_collect(tmp_dir, scm, dvc, run_copy):
     def collect_outs(*args, **kwargs):
         return set(
-            str(out.path_info)
+            str(out)
             for stage in dvc.collect(*args, **kwargs)
             for out in stage.outs
         )

tests/func/test_repro.py

@@ -170,14 +170,13 @@ class TestReproWorkingDirectoryAsOutput(TestDvc):
 
         # NOTE: os.walk() walks in a sorted order and we need dir2 subdirs to
         # be processed before dir1 to load error.dvc first.
-        with patch.object(self.dvc, "collect_stages") as mock_stages:
-            mock_stages.return_value = [
-                nested_stage,
-                Stage.load(self.dvc, error_stage_path),
-            ]
+        self.dvc.stages = [
+            nested_stage,
+            Stage.load(self.dvc, error_stage_path),
+        ]
 
-            with self.assertRaises(StagePathAsOutputError):
-                self.dvc.reproduce(error_stage_path)
+        with self.assertRaises(StagePathAsOutputError):
+            self.dvc.reproduce(error_stage_path)
 
     def test_similar_paths(self):
         # File structure: