From 90a326c7cc371bfbcedb9913184c8ff8904193ff Mon Sep 17 00:00:00 2001 From: Amisha Singla Date: Sun, 1 Jul 2018 23:18:28 +0530 Subject: [PATCH] prints tree to visualise --- osm_fudge/__main__.py | 1 + osm_fudge/index/bk_tree.py | 20 ++++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/osm_fudge/__main__.py b/osm_fudge/__main__.py index 8b6b06b..70e9a4c 100644 --- a/osm_fudge/__main__.py +++ b/osm_fudge/__main__.py @@ -45,6 +45,7 @@ def main(): sys.exit("Error: Empty input file") tree_obj.freeze() + tree_obj.print() if args.query_strings: for query_string in args.query_strings.split(','): gen = tree_obj.nearest(query_string.strip(), max_distance=args.max_distance) diff --git a/osm_fudge/index/bk_tree.py b/osm_fudge/index/bk_tree.py index c14769b..d0cf9d1 100644 --- a/osm_fudge/index/bk_tree.py +++ b/osm_fudge/index/bk_tree.py @@ -6,6 +6,7 @@ from recordclass import recordclass from bisect import bisect_left from bisect import bisect_right +import pydot Node = recordclass('Node', ['value', 'edges', 'children']) @@ -84,3 +85,22 @@ def nearest(self, value, max_distance=None): for index, edge in enumerate(node.edges): if edge >= minimum and edge <= maximum: stack.append(node.children[index]) + + def print(self): + graph = pydot.Dot(graph_type='digraph') + + stack = [] + stack.append(self.root) + + while stack: + node = stack.pop() + if len(node.children): + parent_node = pydot.Node(node.value) + graph.add_node(parent_node) + for index, child in enumerate(node.children): + child_node = pydot.Node(child.value) + graph.add_node(child_node) + graph.add_edge(pydot.Edge(parent_node, child_node, label=node.edges[index])) + stack.append(child) + + graph.write_png('graph.png')