<s> from itertools import count , repeat <EOL> import json <EOL> import networkx as nx <EOL> __author__ = """<STR_LIT>""" <EOL> __all__ = [ '<STR_LIT>' , <EOL> '<STR_LIT>' ] <EOL> def tree_data ( G , root ) : <EOL> """<STR_LIT>""" <EOL> if not G . number_of_nodes ( ) == G . number_of_edges ( ) + <NUM_LIT:1> : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> if not G . is_directed ( ) : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> def add_children ( n , G ) : <EOL> nbrs = G [ n ] <EOL> if len ( nbrs ) == <NUM_LIT:0> : <EOL> return [ ] <EOL> children = [ ] <EOL> for child in nbrs : <EOL> d = dict ( id = child , ** G . node [ child ] ) <EOL> c = add_children ( child , G ) <EOL> if c : <EOL> d [ '<STR_LIT>' ] = c <EOL> children . append ( d ) <EOL> return children <EOL> data = dict ( id = root , ** G . node [ root ] ) <EOL> data [ '<STR_LIT>' ] = add_children ( root , G ) <EOL> return data <EOL> def tree_graph ( data ) : <EOL> """<STR_LIT>""" <EOL> graph = nx . DiGraph ( ) <EOL> def add_children ( parent , children ) : <EOL> for data in children : <EOL> child = data [ '<STR_LIT:id>' ] <EOL> graph . add_edge ( parent , child ) <EOL> grandchildren = data . get ( '<STR_LIT>' , [ ] ) <EOL> if grandchildren : <EOL> add_children ( child , grandchildren ) <EOL> nodedata = dict ( ( str ( k ) , v ) for k , v in data . items ( ) <EOL> if k != '<STR_LIT:id>' and k != '<STR_LIT>' ) <EOL> graph . add_node ( child , attr_dict = nodedata ) <EOL> root = data [ '<STR_LIT:id>' ] <EOL> children = data . get ( '<STR_LIT>' , [ ] ) <EOL> nodedata = dict ( ( k , v ) for k , v in data . items ( ) <EOL> if k != '<STR_LIT:id>' and k != '<STR_LIT>' ) <EOL> graph . add_node ( root , attr_dict = nodedata ) <EOL> add_children ( root , children ) <EOL> return graph </s>
