ENH: pydot_example.py not drawing redundant nodes, building tree form data in...

ENH: pydot_example.py not drawing redundant nodes, building tree form data in mericMeasurement folder, modified colors, shapes, arrows #3

pydot_example.py

+#!/usr/bin/env python3
 import pydot
 from matplotlib import pyplot as pp
 from matplotlib import image as mpimg
@@ -19,9 +20,12 @@ def browse_dir(path):
 		with open('{}/{}'.format(subdir,files[0])) as f:
 			for line in f:
 				line = line.strip()
+				line2 = line.split(';')
 				if line.startswith('# CALLTREE'):
 					print(line)
-					call_list.append(line[11:])
+					tmp = line2[-1].split('_')[0]
+					if tmp not in call_list and tmp != 'init':
+						call_list.append(tmp)
 					print(line.endswith('init_0'))
 					if main_reg is None and line.endswith('init_0'):
 						main_reg = subdir.split('/')[-1]
@@ -29,64 +33,31 @@ def browse_dir(path):
 		reg_structure[subdir.split('/')[-1]] = call_list
 
 	print(reg_structure)
-	#print(main_reg)
-
-tree = {'main': {'load': {'regA': None}, 'compute': {'regA':None, 'regB':None}}}
-
-def walk_dictionaryv2(graph, dictionary, parent_node=None):
-    '''
-    Recursive plotting function for the decision tree stored as a dictionary
-    '''
-
-    for k in dictionary.keys():
-
-        if parent_node is not None:
-
-            from_name = parent_node.get_name().replace("\"", "") + '_' + str(k)
-            from_label = str(k)
-
-            node_from = pydot.Node(from_name, label=from_label)
-            graph.add_node(node_from)
-
-            graph.add_edge( pydot.Edge(parent_node, node_from) )
-
-            if isinstance(dictionary[k], dict): # if interim node
-
-
-                walk_dictionaryv2(graph, dictionary[k], node_from)
-
-            elif dictionary[k] is not None: # if leaf node
-                to_name = str(k) + '_' + str(dictionary[k]) # unique name
-                to_label = str(dictionary[k])
-
-                node_to = pydot.Node(to_name, label=to_label)
-                graph.add_node(node_to)
-
-                graph.add_edge(pydot.Edge(node_from, node_to))
-
-                #node_from.set_name(to_name)
-
-        else:
-
-            from_name =  str(k)
-            from_label = str(k)
-
-            node_from = pydot.Node(from_name, label=from_label)
-            walk_dictionaryv2(graph, dictionary[k], node_from)
-
-
-def plot_tree(tree, name):
-
-    # first you create a new graph, you do that with pydot.Dot()
-    graph = pydot.Dot(graph_type='graph',dpi=300)
-
-    walk_dictionaryv2(graph, tree)
-    graph.write_png(name+'.png')
+	G = pydot.Dot(graph_type='digraph',dpi=800)
+	for k in reg_structure.keys():
+		if k == main_reg:
+			n = pydot.Node(k, style = '"rounded,filled"', fillcolor = '#ccaaff', shape = 'rectangle')
+		else:
+		    n = pydot.Node(k, style = 'filled', fillcolor = '#d5e4f5', shape = 'square')
+		G.add_node(n)
+
+	for k,v in reg_structure.items():
+		for e in v:
+			G.add_edge(pydot.Edge(e,k, arrowsize = 0.5))
+
+	G.write_png('pokus.png')
+	tree_img = mpimg.imread('pokus.png')
+	pp.imshow(tree_img)
+	pp.axis('off')
+	pp.show()
 
+	#print(main_reg)
 
-browse_dir('/home/david/Desktop/mericMeasurement')
-plot_tree(tree,'test_tree')
-tree_img = mpimg.imread('test_tree.png')
-pp.imshow(tree_img)
-pp.axis('off')
-pp.show()
+if __name__=='__main__':
+    browse_dir('/home/david/Desktop/mericMeasurement')
+    #plot_tree(tree,'test_tree')
+    #plot_test('test_tree')
+    #tree_img = mpimg.imread('test_tree.png')
+    #pp.imshow(tree_img)
+    #pp.axis('off')
+    #pp.show()