[engine] Move the visualizer into LocalScheduler.

- Migrate the meat of the visualizer into `LocalScheduler.visualize_graph`.
- Add basic tests for coverage.

Testing Done:
https://travis-ci.org/pantsbuild/pants/builds/120774094

+ local testing against viz and viz-fs.

Bugs closed: 3138, 3141

Reviewed at https://rbcommons.com/s/twitter/r/3649/

src/python/pants/engine/exp/examples/visualizer.py

@@ -15,80 +15,16 @@
 from pants.engine.exp.engine import LocalSerialEngine
 from pants.engine.exp.examples.planners import setup_json_scheduler
 from pants.engine.exp.fs import PathGlobs
-from pants.engine.exp.nodes import Noop, SelectNode, TaskNode, Throw
-from pants.util.contextutil import temporary_file, temporary_file_path
-
-
-def format_type(node):
-  if type(node) == TaskNode:
-    return node.func.__name__
-  return type(node).__name__
-
-
-def format_subject(node):
-  if node.variants:
-    return '({})@{}'.format(node.subject, ','.join('{}={}'.format(k, v) for k, v in node.variants))
-  else:
-    return '({})'.format(node.subject)
-
-
-def format_product(node):
-  if type(node) == SelectNode and node.variant_key:
-    return '{}@{}'.format(node.product.__name__, node.variant_key)
-  return node.product.__name__
-
-
-def format_node(node, state):
-  return '{}:{}:{} == {}'.format(format_product(node),
-                                 format_subject(node),
-                                 format_type(node),
-                                 str(state).replace('"', '\\"'))
-
-
-# NB: there are only 12 colors in `set312`.
-colorscheme = 'set312'
-max_colors = 12
-colors = {}
-
-
-def format_color(node, node_state):
-  if type(node_state) is Throw:
-    return 'tomato'
-  elif type(node_state) is Noop:
-    return 'white'
-  key = node.product
-  return colors.setdefault(key, (len(colors) % max_colors) + 1)
-
-
-def create_digraph(scheduler, storage, request):
-
-  yield 'digraph plans {'
-  yield '  node[colorscheme={}];'.format(colorscheme)
-  yield '  concentrate=true;'
-  yield '  rankdir=LR;'
-
-  for ((node, node_state), dependency_entries) in scheduler.product_graph.walk(request.roots):
-    node_str = format_node(node, node_state)
-
-    yield (' "{node}" [style=filled, fillcolor={color}];'
-            .format(color=format_color(node, node_state),
-                    node=node_str))
-
-    for (dep, dep_state) in dependency_entries:
-      yield '  "{}" -> "{}"'.format(node_str, format_node(dep, dep_state))
-
-  yield '}'
+from pants.util.contextutil import temporary_file_path
 
 
 def visualize_execution_graph(scheduler, storage, request):
-  with temporary_file(cleanup=False, suffix='.dot') as fp:
-    for line in create_digraph(scheduler, storage, request):
-      fp.write(line)
-      fp.write('\n')
+  with temporary_file_path(cleanup=False, suffix='.dot') as dot_file:
+    scheduler.visualize_graph_to_file(request.roots, dot_file)
+    print('dot file saved to: {}'.format(dot_file))
 
-  print('dot file saved to: {}'.format(fp.name))
   with temporary_file_path(cleanup=False, suffix='.svg') as image_file:
-    subprocess.check_call('dot -Tsvg -o{} {}'.format(image_file, fp.name), shell=True)
+    subprocess.check_call('dot -Tsvg -o{} {}'.format(image_file, dot_file), shell=True)
     print('svg file saved to: {}'.format(image_file))
     binary_util.ui_open(image_file)
 
@@ -119,7 +55,6 @@ def usage(error_message):
 
   build_root = args.pop(0)
 
-
   if not os.path.isdir(build_root):
     usage('First argument must be a valid build root, {} is not a directory.'.format(build_root))
   build_root = os.path.realpath(build_root)

src/python/pants/engine/exp/scheduler.py

@@ -242,6 +242,58 @@ def _walk(entries):
     for entry in _walk(_filtered_entries(roots)):
       yield entry
 
+  def visualize(self, roots):
+    """Visualize a graph walk by generating graphviz `dot` output.
+
+    :param iterable roots: An iterable of the root nodes to begin the graph walk from.
+    """
+    viz_colors = {}
+    viz_color_scheme = 'set312'  # NB: There are only 12 colors in `set312`.
+    viz_max_colors = 12
+
+    def format_color(node, node_state):
+      if type(node_state) is Throw:
+        return 'tomato'
+      elif type(node_state) is Noop:
+        return 'white'
+      return viz_colors.setdefault(node.product, (len(viz_colors) % viz_max_colors) + 1)
+
+    def format_type(node):
+      return node.func.__name__ if type(node) is TaskNode else type(node).__name__
+
+    def format_subject(node):
+      if node.variants:
+        return '({})@{}'.format(node.subject,
+                                ','.join('{}={}'.format(k, v) for k, v in node.variants))
+      else:
+        return '({})'.format(node.subject)
+
+    def format_product(node):
+      if type(node) is SelectNode and node.variant_key:
+        return '{}@{}'.format(node.product.__name__, node.variant_key)
+      return node.product.__name__
+
+    def format_node(node, state):
+      return '{}:{}:{} == {}'.format(format_product(node),
+                                     format_subject(node),
+                                     format_type(node),
+                                     str(state).replace('"', '\\"'))
+
+    yield 'digraph plans {'
+    yield '  node[colorscheme={}];'.format(viz_color_scheme)
+    yield '  concentrate=true;'
+    yield '  rankdir=LR;'
+
+    for ((node, node_state), dependency_entries) in self.walk(roots):
+      node_str = format_node(node, node_state)
+
+      yield ' "{}" [style=filled, fillcolor={}];'.format(node_str, format_color(node, node_state))
+
+      for (dep, dep_state) in dependency_entries:
+        yield '  "{}" -> "{}"'.format(node_str, format_node(dep, dep_state))
+
+    yield '}'
+
 
 class ExecutionRequest(datatype('ExecutionRequest', ['roots'])):
   """Holds the roots for an execution, which might have been requested by a user.
@@ -386,6 +438,17 @@ def __init__(self, goals, tasks, storage, project_tree, graph_lock=None, graph_v
     self._product_graph_lock = graph_lock or threading.RLock()
     self._step_id = 0
 
+  def visualize_graph_to_file(self, roots, filename):
+    """Visualize a graph walk by writing graphviz `dot` output to a file.
+
+    :param iterable roots: An iterable of the root nodes to begin the graph walk from.
+    :param str filename: The filename to output the graphviz output to.
+    """
+    with self._product_graph_lock, open(filename, 'wb') as fh:
+      for line in self.product_graph.visualize(roots):
+        fh.write(line)
+        fh.write('\n')
+
   def _create_step(self, node):
     """Creates a Step and Promise with the currently available dependencies of the given Node.
 

tests/python/pants_test/engine/exp/test_scheduler.py

@@ -20,6 +20,7 @@
 from pants.engine.exp.nodes import (ConflictingProducersError, DependenciesNode, Return, SelectNode,
                                     Throw, Waiting)
 from pants.engine.exp.scheduler import ProductGraph
+from pants.util.contextutil import temporary_dir
 
 
 class SchedulerTest(unittest.TestCase):
@@ -244,6 +245,23 @@ def test_sibling_specs(self):
     # And that an subdirectory address is not.
     self.assertNotIn(self.managed_guava, root_value)
 
+  def test_scheduler_visualize(self):
+    spec = self.spec_parser.parse_spec('3rdparty/jvm:')
+    build_request = self.request_specs(['list'], spec)
+    self.build_and_walk(build_request)
+
+    graphviz_output = '\n'.join(self.scheduler.product_graph.visualize(build_request.roots))
+
+    with temporary_dir() as td:
+      output_path = os.path.join(td, 'output.dot')
+      self.scheduler.visualize_graph_to_file(build_request.roots, output_path)
+      with open(output_path, 'rb') as fh:
+        graphviz_disk_output = fh.read().strip()
+
+    self.assertEqual(graphviz_output, graphviz_disk_output)
+    self.assertIn('digraph', graphviz_output)
+    self.assertIn(' -> ', graphviz_output)
+
 
 # TODO: Expand test coverage here.
 class ProductGraphTest(unittest.TestCase):