Merge 2e75cbc into 63b0e8b

aiida/backends/tests/cmdline/commands/test_graph.py

@@ -69,7 +69,7 @@ def test_generate_graph(self):
         """
         # Get a PK of a node which exists
         root_node = str(self.node.pk)
-        filename = root_node + '.dot'
+        filename = root_node + '.dot.pdf'
         options = [root_node]
         try:
             result = self.cli_runner.invoke(cmd_graph.generate, options)
@@ -89,7 +89,7 @@ def test_catch_bad_pk(self):
         # Forbidden pk
         for root_node in ['xyz', '-5', '3.14']:
             options = [root_node]
-            filename = root_node + '.dot'
+            filename = root_node + '.dot.pdf'
             try:
                 result = self.cli_runner.invoke(cmd_graph.generate, options)
                 self.assertIsNotNone(result.exception)
@@ -108,7 +108,7 @@ def test_catch_bad_pk(self):
             pass
         #  Make sure verdi graph rejects this non-existant pk
         try:
-            filename = str(root_node) + '.dot'
+            filename = str(root_node) + '.dot.pdf'
             options = [str(root_node)]
             result = self.cli_runner.invoke(cmd_graph.generate, options)
             self.assertIsNotNone(result.exception)
@@ -123,7 +123,7 @@ def test_check_recursion_flags(self):
         positive ints
         """
         root_node = str(self.node.pk)
-        filename = root_node + '.dot'
+        filename = root_node + '.dot.pdf'
 
         # Test that the options don't fail
         for opt in ['-a', '--ancestor-depth', '-d', '--descendant-depth']:
@@ -162,9 +162,9 @@ def test_check_io_flags(self):
         Test the input and output flags work.
         """
         root_node = str(self.node.pk)
-        filename = root_node + '.dot'
+        filename = root_node + '.dot.pdf'
 
-        for flag in ['-i', '--inputs', '-o', '--outputs']:
+        for flag in ['-i', '--process-in', '-o', '--process-out']:
             options = [flag, root_node]
             try:
                 result = self.cli_runner.invoke(cmd_graph.generate, options)
@@ -185,8 +185,8 @@ def test_output_format(self):
             # supported on a given OS (printed by '$ dot -T?') but here
             # we just use the built-ins dot and canon as a minimal check that
             # the option works. After all, this test is for the cmdline.
-            for fileformat in ['dot', 'canon']:
-                filename = root_node + '.' + fileformat
+            for fileformat in ['pdf', 'png']:
+                filename = root_node + '.dot.' + fileformat
                 options = [option, fileformat, root_node]
                 try:
                     result = self.cli_runner.invoke(cmd_graph.generate, options)

aiida/cmdline/commands/cmd_graph.py

@@ -23,12 +23,18 @@
 def verdi_graph():
     """
     Create visual representations of part of the provenance graph.
-    Requires that `graphviz<https://graphviz.org/download>` be installed.
+    Requires that `python-graphviz<https://graphviz.readthedocs.io>` be installed.
     """
 
 
 @verdi_graph.command('generate')
 @arguments.NODE('root_node')
+@click.option(
+    "-l",
+    "--link-types",
+    help="The link types to include: 'data' includes only 'input_calc' and 'create' links (data provenance only),  'logic' includes only 'input_work' and 'return' links (logical provenance only).",
+    default="all",
+    type=click.Choice(['all', 'data', 'logic']))
 @click.option(
     '-a',
     '--ancestor-depth',
@@ -37,33 +43,49 @@ def verdi_graph():
 @click.option(
     '-d',
     '--descendant-depth',
-    help='The maximum depth when recursing through the descendants, if not set it will recurse to the end',
+    help='The maximum depth when recursing through the descendants. If not set it will recurse to the end.',
     type=click.IntRange(min=0))
-@click.option('-o', '--outputs', is_flag=True, help='Always show all outputs of a calculation.')
-@click.option('-i', '--inputs', is_flag=True, help='Always show all inputs of a calculation.')
+@click.option('-o', '--process-out', is_flag=True, help='Show outgoing links for all processes.')
+@click.option('-i', '--process-in', is_flag=True, help='Show incoming links for all processes.')
+@click.option('-p', '--print-info', is_flag=True, help="Print verbose information of the graph traversal.")
 @click.option(
-    '-f',
-    '--output-format',
-    help='The output format, something that can be recognized by graphviz'
-    '(see http://www.graphviz.org/doc/info/output.html).',
+    '-e',
+    '--engine',
+    help="The graphviz engine, e.g. 'dot', 'circo', ... "
+    "(see http://www.graphviz.org/doc/info/output.html)",
     default='dot')
+@click.option('-f', '--output-format', help="The output format used for rendering ('pdf', 'png', etc.).", default='pdf')
+@click.option('-v', '--view', is_flag=True, help="Open the rendered result with the default application.")
 @decorators.with_dbenv()
-def generate(root_node, ancestor_depth, descendant_depth, outputs, inputs, output_format):
-    """Generate a graph for a given ROOT_NODE."""
-    from aiida.tools.visualization.graphviz import draw_graph
+def generate(root_node, link_types, ancestor_depth, descendant_depth, process_out, process_in, engine, print_info,
+             output_format, view):
+    """
+    Generate a graph from a ROOT_NODE (specified by pk or uuid).
+    """
+    # pylint: disable=too-many-arguments
+    from aiida.tools.visualization.graph import Graph
+    print_func = echo.echo_info if print_info else None
+    link_types = {"all": (), "logic": ("input_work", "return"), "data": ("input_calc", "create")}[link_types]
 
-    try:
-        exit_status, output_file_name = draw_graph(
-            root_node,
-            ancestor_depth=ancestor_depth,
-            descendant_depth=descendant_depth,
-            image_format=output_format,
-            include_calculation_inputs=inputs,
-            include_calculation_outputs=outputs)
-    except OSError as exception:
-        echo.echo_critical(str(exception))
+    echo.echo_info("Initiating graphviz engine: {}".format(engine))
+    graph = Graph(engine=engine)
+    echo.echo_info("Recursing ancestors, max depth={}".format(ancestor_depth))
+    graph.recurse_ancestors(
+        root_node,
+        depth=ancestor_depth,
+        link_types=link_types,
+        annotate_links="both",
+        include_process_outputs=process_out,
+        print_func=print_func)
+    echo.echo_info("Recursing descendants, max depth={}".format(descendant_depth))
+    graph.recurse_descendants(
+        root_node,
+        depth=descendant_depth,
+        link_types=link_types,
+        annotate_links="both",
+        include_process_inputs=process_in,
+        print_func=print_func)
+    output_file_name = graph.graphviz.render(
+        filename='{}.{}'.format(root_node.pk, engine), format=output_format, view=view, cleanup=True)
 
-    if exit_status:
-        echo.echo_critical('Failed to generate graph')
-    else:
-        echo.echo_success('Output file is {}'.format(output_file_name))
+    echo.echo_success("Output file: {}".format(output_file_name))