From b707dbd0fe405ce296198b5e1b6d4e673a9a11f1 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Mon, 21 Aug 2017 22:59:58 +0300
Subject: [PATCH] dvc: add 'dvc visual'

This command generates dependency graph for specified target,
saves it in .dot file and also draws it into .svg file.

Fixes #39

Signed-off-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/cli.py            | 12 ++++++
 dvc/command/visual.py | 90 +++++++++++++++++++++++++++++++++++++++++++
 requirements.txt      |  1 +
 setup.py              |  1 +
 4 files changed, 104 insertions(+)
 create mode 100644 dvc/command/visual.py

diff --git a/dvc/cli.py b/dvc/cli.py
index 4dcc4ce80..daf220e7e 100644
--- a/dvc/cli.py
+++ b/dvc/cli.py
@@ -14,6 +14,7 @@ from dvc.command.import_file import CmdImportFile
 from dvc.command.target import CmdTarget
 from dvc.command.instance_create import CmdInstanceCreate
 from dvc.command.config import CmdConfig
+from dvc.command.visual import CmdVisual
 from dvc.command.test import CmdTest
 from dvc.config import Config
 from dvc import VERSION
@@ -388,6 +389,17 @@ def parse_args(argv=None):
                         help='Option value')
     config_parser.set_defaults(func=CmdConfig)
 
+    # Visual
+    vz_parser = subparsers.add_parser(
+                        'visual',
+                        parents=[parent_parser],
+                        help='Create a dependency graph for data')
+    vz_parser.add_argument(
+                        'target',
+                        nargs='*',
+                        help='Target data')
+    vz_parser.set_defaults(func=CmdVisual)
+
     if isinstance(argv, str):
         argv = argv.split()
 
diff --git a/dvc/command/visual.py b/dvc/command/visual.py
new file mode 100644
index 000000000..41abcc3d2
--- /dev/null
+++ b/dvc/command/visual.py
@@ -0,0 +1,90 @@
+import os
+import networkx as nx
+
+from dvc.command.traverse import Traverse
+from dvc.command.init import CmdInit
+
+from dvc.logger import Logger
+from dvc.state_file import StateFile
+
+class CmdVisual(Traverse):
+    def __init__(self, settings):
+        super(CmdVisual, self).__init__(settings, "collect", do_not_start_from_root=False)
+        self.g = nx.DiGraph()
+        self.subs = []
+
+    def draw(self, g, target, fname):
+        try:
+            A = nx.nx_agraph.to_agraph(g)
+            A.write(fname + '.dot')
+            A.draw(fname + '.svg', prog='dot')
+        except Exception as exc:
+            Logger.error('Failed to draw dependency graph for {}: {}'.format(target, exc))
+            return 1
+
+        return 0
+
+    def find_sub(self, fname):
+        for s in self.subs:
+            if fname in s.nodes():
+                return s
+
+        return self.g
+
+    def draw_targets(self):
+        if self.parsed_args.target == None:
+            target = self.settings.config.data_dir
+            return self.draw(self.g, target, target)
+
+        for t in self.parsed_args.target:
+            fname = os.path.basename(os.path.normpath(t))
+            s = self.find_sub(t)
+
+            ret = self.draw(s, t, fname)
+            if ret != 0:
+                return ret
+
+        return 0
+
+    def run(self):
+        saved_targets = self.settings.parsed_args.target
+        self.settings.parsed_args.target = [self.settings.config.data_dir]
+ 
+        ret = super(CmdVisual, self).run()
+
+        self.settings.parsed_args.target = saved_targets
+
+        if ret != 0:
+            Logger.error('Failed to build dependency graph for the project')
+            return 1
+
+        # Try to find independent clusters which might occure
+        # when a bunch of data items were used independently.
+        self.subs = nx.weakly_connected_component_subgraphs(self.g)
+
+        return self.draw_targets()
+
+    def process_file(self, target):
+        data_item = self._get_data_item(target)
+        name = data_item.data.relative
+        state = StateFile.load(data_item.state.relative, self.git)
+
+        self.g.add_node(name)
+
+        for i in state.input_files:
+            self.g.add_node(i)
+            self.g.add_edge(i, name)
+
+        for o in state.output_files:
+            if o == name:
+                continue
+            self.g.add_node(o)
+            self.g.add_edge(name, o)
+
+    @property
+    def no_git_actions(self):
+        return True
+
+    @staticmethod
+    def not_committed_changes_warning():
+        pass
diff --git a/requirements.txt b/requirements.txt
index 8e3d4ad16..8754dcc29 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -87,3 +87,4 @@ wcwidth==0.1.7
 colorama==0.3.9
 google-compute-engine==2.4.1
 configobj==5.0.6
+networkx==1.11
diff --git a/setup.py b/setup.py
index d47341ba2..28868597e 100644
--- a/setup.py
+++ b/setup.py
@@ -93,6 +93,7 @@ install_requires = [
     'colorama==0.3.9',
     'google-compute-engine==2.4.1',
     'configobj==5.0.6',
+    'networkx==1.11',
 ]
 
 if platform.system() == 'Darwin':