Merge "Enable conversion of the tree nodes into a digraph"

This commit is contained in:
Jenkins 2015-11-18 22:58:32 +00:00 committed by Gerrit Code Review
commit ed88aa4c39
2 changed files with 39 additions and 0 deletions
taskflow
tests/unit
types

@ -511,6 +511,27 @@ CEO
self.assertEqual(['mammal', 'reptile', 'horse',
'primate', 'monkey', 'human'], things)
def test_to_diagraph(self):
root = self._make_species()
g = root.to_digraph()
self.assertEqual(root.child_count(only_direct=False) + 1, len(g))
for node in root.dfs_iter(include_self=True):
self.assertIn(node.item, g)
self.assertEqual([], g.predecessors('animal'))
self.assertEqual(['animal'], g.predecessors('reptile'))
self.assertEqual(['primate'], g.predecessors('human'))
self.assertEqual(['mammal'], g.predecessors('primate'))
self.assertEqual(['animal'], g.predecessors('mammal'))
self.assertEqual(['mammal', 'reptile'], g.successors('animal'))
def test_to_digraph_retains_metadata(self):
root = tree.Node("chickens", alive=True)
dead_chicken = tree.Node("chicken.1", alive=False)
root.add(dead_chicken)
g = root.to_digraph()
self.assertEqual(g.node['chickens'], {'alive': True})
self.assertEqual(g.node['chicken.1'], {'alive': False})
class OrderedSetTest(test.TestCase):

@ -22,6 +22,7 @@ import os
import six
from taskflow.types import graph
from taskflow.utils import iter_utils
from taskflow.utils import misc
@ -390,3 +391,20 @@ class Node(object):
return _BFSIter(self,
include_self=include_self,
right_to_left=right_to_left)
def to_digraph(self):
"""Converts this node + its children into a ordered directed graph.
The graph returned will have the same structure as the
this node and its children (and tree node metadata will be translated
into graph node metadata).
:returns: a directed graph
:rtype: :py:class:`taskflow.types.graph.OrderedDiGraph`
"""
g = graph.OrderedDiGraph()
for node in self.bfs_iter(include_self=True, right_to_left=True):
g.add_node(node.item, attr_dict=node.metadata)
if node is not self:
g.add_edge(node.parent.item, node.item)
return g