Fixing the def main function, properly calling it from treecutter.
[treecutter.git] / treecutter / trie.py
1 #!/usr/bin/python
2 import pygraphviz as pgv
3
4 class Node():
5     def __init__(self,token,value):
6         self._token = token
7         self._value = value
8         self._children = []
9
10     def token(self):
11         return self._token
12
13     def value(self):
14         return self._value
15
16     def children(self):
17         return self._children
18
19 class Trie():
20     def __init__(self):
21         self._root = []
22
23     def __iter__(self):
24         return self.inorder(self._root)
25
26     def inorder(self,t):
27         for l in t:
28             yield l.value()
29             for x in self.inorder(l.children()):
30                 yield x
31
32     def _add(self,trie, key, content):
33         # is the key a leaf
34         k = key.pop(0)
35         if key == []:
36             node = Node(k,content)
37             trie.append(node)
38         else:
39             for ch in trie:
40                 if ch.token() == k:
41                     self._add(ch.children(), key, content)
42
43     def add(self,key, content):
44         self._add(self._root, key, content)
45
46     def _graph(self, trie, G):
47         for l in trie:
48             G.add_node(l.token())
49             for ch in l.children():
50                 G.add_edge(l.token(),ch.token())
51                 self._graph(l.children(), G)
52
53     def graph(self):
54         G = pgv.AGraph(directed=True)
55         G.add_node("sitemap")
56         for ch in self._root:
57             G.add_edge("sitemap",ch.token())
58         self._graph(self._root, G)
59 #        G.layout('dot')
60 #        G.draw('g.png')
61 #        print G.string()
62
63     def _menu(self, trie, lang, page, css):
64         html = "<ul%s>\n" % css
65         for l in trie:
66             sel = ''
67             p = l.value().page(lang)
68             if p == page:
69                 sel = ' class="selected"'
70             if p != None:
71                 html += '<li%s><a href="%s">%s</a>\n' \
72                     % (sel,l.value().link(),p.menu())
73             else:
74                 html += '<li%s><a href="%s.en" hreflang="en">%s</a>*\n' \
75                     % (sel,l.value().link(), l.value().page('en').menu())
76             if l.children():
77                 html += self._menu(l.children(), lang, page, "")
78         html += "</ul>\n"
79         return html
80
81     def menu(self,lang,page,cssclass):
82         css = ''
83         if cssclass:
84             css = ' class="'+cssclass+'"'
85         return self._menu(self._root, lang, page, css)