| 1 | # Copyright (c) 2003-2006 LOGILAB S.A. (Paris, FRANCE). |
|---|
| 2 | # http://www.logilab.fr/ -- mailto:contact@logilab.fr |
|---|
| 3 | # |
|---|
| 4 | # This program is free software; you can redistribute it and/or modify it under |
|---|
| 5 | # the terms of the GNU General Public License as published by the Free Software |
|---|
| 6 | # Foundation; either version 2 of the License, or (at your option) any later |
|---|
| 7 | # version. |
|---|
| 8 | # |
|---|
| 9 | # This program is distributed in the hope that it will be useful, but WITHOUT |
|---|
| 10 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS |
|---|
| 11 | # FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. |
|---|
| 12 | # |
|---|
| 13 | # You should have received a copy of the GNU General Public License along with |
|---|
| 14 | # this program; if not, write to the Free Software Foundation, Inc., |
|---|
| 15 | # 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA. |
|---|
| 16 | """ |
|---|
| 17 | base class to represent tree structure |
|---|
| 18 | """ |
|---|
| 19 | |
|---|
| 20 | import sys |
|---|
| 21 | |
|---|
| 22 | from logilab.common import flatten |
|---|
| 23 | from logilab.common.visitor import VisitedMixIn, FilteredIterator, no_filter |
|---|
| 24 | |
|---|
| 25 | ## Exceptions ################################################################# |
|---|
| 26 | |
|---|
| 27 | class NodeNotFound(Exception): |
|---|
| 28 | """raised when a node has not been found""" |
|---|
| 29 | |
|---|
| 30 | EX_SIBLING_NOT_FOUND = "No such sibling as '%s'" |
|---|
| 31 | EX_CHILD_NOT_FOUND = "No such child as '%s'" |
|---|
| 32 | EX_NODE_NOT_FOUND = "No such node as '%s'" |
|---|
| 33 | |
|---|
| 34 | |
|---|
| 35 | # Base node ################################################################### |
|---|
| 36 | |
|---|
| 37 | class Node(object): |
|---|
| 38 | """a basic tree node, caracterised by an id""" |
|---|
| 39 | |
|---|
| 40 | def __init__(self, nid=None) : |
|---|
| 41 | self.id = nid |
|---|
| 42 | # navigation |
|---|
| 43 | self.parent = None |
|---|
| 44 | self.children = [] |
|---|
| 45 | |
|---|
| 46 | def __str__(self, indent=0): |
|---|
| 47 | s = ['%s%s %s' % (' '*indent, self.__class__.__name__, self.id)] |
|---|
| 48 | indent += 2 |
|---|
| 49 | for child in self.children: |
|---|
| 50 | try: |
|---|
| 51 | s.append(child.__str__(indent)) |
|---|
| 52 | except TypeError: |
|---|
| 53 | s.append(child.__str__()) |
|---|
| 54 | return '\n'.join(s) |
|---|
| 55 | |
|---|
| 56 | |
|---|
| 57 | def is_leaf(self): |
|---|
| 58 | return not self.children |
|---|
| 59 | |
|---|
| 60 | def append(self, child): |
|---|
| 61 | """add a node to children""" |
|---|
| 62 | self.children.append(child) |
|---|
| 63 | child.parent = self |
|---|
| 64 | |
|---|
| 65 | def remove(self, child): |
|---|
| 66 | """remove a child node""" |
|---|
| 67 | self.children.remove(child) |
|---|
| 68 | child.parent = None |
|---|
| 69 | |
|---|
| 70 | def insert(self, index, child): |
|---|
| 71 | """insert a child node""" |
|---|
| 72 | self.children.insert(index, child) |
|---|
| 73 | child.parent = self |
|---|
| 74 | |
|---|
| 75 | def replace(self, old_child, new_child): |
|---|
| 76 | """replace a child node with another""" |
|---|
| 77 | i = self.children.index(old_child) |
|---|
| 78 | self.children.pop(i) |
|---|
| 79 | self.children.insert(i, new_child) |
|---|
| 80 | new_child.parent = self |
|---|
| 81 | |
|---|
| 82 | def get_sibling(self, nid): |
|---|
| 83 | """return the sibling node that has given id""" |
|---|
| 84 | try: |
|---|
| 85 | return self.parent.get_child_by_id(nid) |
|---|
| 86 | except NodeNotFound : |
|---|
| 87 | raise NodeNotFound(EX_SIBLING_NOT_FOUND % nid) |
|---|
| 88 | |
|---|
| 89 | def next_sibling(self): |
|---|
| 90 | """ |
|---|
| 91 | return the next sibling for this node if any |
|---|
| 92 | """ |
|---|
| 93 | parent = self.parent |
|---|
| 94 | if parent is None: |
|---|
| 95 | # root node has no sibling |
|---|
| 96 | return None |
|---|
| 97 | index = parent.children.index(self) |
|---|
| 98 | try: |
|---|
| 99 | return parent.children[index+1] |
|---|
| 100 | except IndexError: |
|---|
| 101 | return None |
|---|
| 102 | |
|---|
| 103 | def previous_sibling(self): |
|---|
| 104 | """ |
|---|
| 105 | return the previous sibling for this node if any |
|---|
| 106 | """ |
|---|
| 107 | parent = self.parent |
|---|
| 108 | if parent is None: |
|---|
| 109 | # root node has no sibling |
|---|
| 110 | return None |
|---|
| 111 | index = parent.children.index(self) |
|---|
| 112 | if index > 0: |
|---|
| 113 | return parent.children[index-1] |
|---|
| 114 | return None |
|---|
| 115 | |
|---|
| 116 | def get_node_by_id(self, nid): |
|---|
| 117 | """ |
|---|
| 118 | return node in whole hierarchy that has given id |
|---|
| 119 | """ |
|---|
| 120 | root = self.root() |
|---|
| 121 | try: |
|---|
| 122 | return root.get_child_by_id(nid, 1) |
|---|
| 123 | except NodeNotFound : |
|---|
| 124 | raise NodeNotFound(EX_NODE_NOT_FOUND % nid) |
|---|
| 125 | |
|---|
| 126 | def get_child_by_id(self, nid, recurse=None): |
|---|
| 127 | """ |
|---|
| 128 | return child of given id |
|---|
| 129 | """ |
|---|
| 130 | if self.id == nid: |
|---|
| 131 | return self |
|---|
| 132 | for c in self.children : |
|---|
| 133 | if recurse: |
|---|
| 134 | try: |
|---|
| 135 | return c.get_child_by_id(nid, 1) |
|---|
| 136 | except NodeNotFound : |
|---|
| 137 | continue |
|---|
| 138 | if c.id == nid : |
|---|
| 139 | return c |
|---|
| 140 | raise NodeNotFound(EX_CHILD_NOT_FOUND % nid) |
|---|
| 141 | |
|---|
| 142 | def get_child_by_path(self, path): |
|---|
| 143 | """ |
|---|
| 144 | return child of given path (path is a list of ids) |
|---|
| 145 | """ |
|---|
| 146 | if len(path) > 0 and path[0] == self.id: |
|---|
| 147 | if len(path) == 1 : |
|---|
| 148 | return self |
|---|
| 149 | else : |
|---|
| 150 | for c in self.children : |
|---|
| 151 | try: |
|---|
| 152 | return c.get_child_by_path(path[1:]) |
|---|
| 153 | except NodeNotFound : |
|---|
| 154 | pass |
|---|
| 155 | raise NodeNotFound(EX_CHILD_NOT_FOUND % path) |
|---|
| 156 | |
|---|
| 157 | def depth(self): |
|---|
| 158 | """ |
|---|
| 159 | return depth of this node in the tree |
|---|
| 160 | """ |
|---|
| 161 | if self.parent is not None: |
|---|
| 162 | return 1 + self.parent.depth() |
|---|
| 163 | else : |
|---|
| 164 | return 0 |
|---|
| 165 | |
|---|
| 166 | def depth_down(self): |
|---|
| 167 | """ |
|---|
| 168 | return depth of the tree from this node |
|---|
| 169 | """ |
|---|
| 170 | if self.children: |
|---|
| 171 | return 1 + max([c.depth_down() for c in self.children]) |
|---|
| 172 | return 1 |
|---|
| 173 | |
|---|
| 174 | def width(self): |
|---|
| 175 | """ |
|---|
| 176 | return the width of the tree from this node |
|---|
| 177 | """ |
|---|
| 178 | return len(self.leaves()) |
|---|
| 179 | |
|---|
| 180 | def root(self): |
|---|
| 181 | """ |
|---|
| 182 | return the root node of the tree |
|---|
| 183 | """ |
|---|
| 184 | if self.parent is not None: |
|---|
| 185 | return self.parent.root() |
|---|
| 186 | return self |
|---|
| 187 | |
|---|
| 188 | def leaves(self): |
|---|
| 189 | """ |
|---|
| 190 | return a list with all the leaves nodes descendant from this node |
|---|
| 191 | """ |
|---|
| 192 | leaves = [] |
|---|
| 193 | if self.children: |
|---|
| 194 | for child in self.children: |
|---|
| 195 | leaves += child.leaves() |
|---|
| 196 | return leaves |
|---|
| 197 | else: |
|---|
| 198 | return [self] |
|---|
| 199 | |
|---|
| 200 | def __iter__(self): |
|---|
| 201 | return iter(self.children) |
|---|
| 202 | |
|---|
| 203 | def flatten(self, _list=None): |
|---|
| 204 | """ |
|---|
| 205 | return a list with all the nodes descendant from this node |
|---|
| 206 | """ |
|---|
| 207 | if _list is None: |
|---|
| 208 | _list = [] |
|---|
| 209 | _list.append(self) |
|---|
| 210 | for c in self.children: |
|---|
| 211 | c.flatten(_list) |
|---|
| 212 | return _list |
|---|
| 213 | |
|---|
| 214 | def lineage(self): |
|---|
| 215 | """ |
|---|
| 216 | return list of parents up to root node |
|---|
| 217 | """ |
|---|
| 218 | lst = [self] |
|---|
| 219 | if self.parent is not None: |
|---|
| 220 | lst.extend(self.parent.lineage()) |
|---|
| 221 | return lst |
|---|
| 222 | |
|---|
| 223 | class VNode(Node, VisitedMixIn): |
|---|
| 224 | """a visitable node |
|---|
| 225 | """ |
|---|
| 226 | pass |
|---|
| 227 | |
|---|
| 228 | |
|---|
| 229 | class BinaryNode(VNode): |
|---|
| 230 | """a binary node (ie only two children |
|---|
| 231 | """ |
|---|
| 232 | def __init__(self, lhs=None, rhs=None) : |
|---|
| 233 | VNode.__init__(self) |
|---|
| 234 | if lhs is not None or rhs is not None: |
|---|
| 235 | assert lhs and rhs |
|---|
| 236 | self.append(lhs) |
|---|
| 237 | self.append(rhs) |
|---|
| 238 | |
|---|
| 239 | def remove(self, child): |
|---|
| 240 | """remove the child and replace this node with the other child |
|---|
| 241 | """ |
|---|
| 242 | self.children.remove(child) |
|---|
| 243 | self.parent.replace(self, self.children[0]) |
|---|
| 244 | |
|---|
| 245 | def get_parts(self): |
|---|
| 246 | """ |
|---|
| 247 | return the left hand side and the right hand side of this node |
|---|
| 248 | """ |
|---|
| 249 | return self.children[0], self.children[1] |
|---|
| 250 | |
|---|
| 251 | |
|---|
| 252 | |
|---|
| 253 | if sys.version_info[0:2] >= (2, 2): |
|---|
| 254 | list_class = list |
|---|
| 255 | else: |
|---|
| 256 | from UserList import UserList |
|---|
| 257 | list_class = UserList |
|---|
| 258 | |
|---|
| 259 | class ListNode(VNode, list_class): |
|---|
| 260 | """Used to manipulate Nodes as Lists |
|---|
| 261 | """ |
|---|
| 262 | def __init__(self): |
|---|
| 263 | list_class.__init__(self) |
|---|
| 264 | VNode.__init__(self) |
|---|
| 265 | self.children = self |
|---|
| 266 | |
|---|
| 267 | def __str__(self, indent=0): |
|---|
| 268 | return '%s%s %s' % (indent*' ', self.__class__.__name__, |
|---|
| 269 | ', '.join([str(v) for v in self])) |
|---|
| 270 | |
|---|
| 271 | def append(self, child): |
|---|
| 272 | """add a node to children""" |
|---|
| 273 | list_class.append(self, child) |
|---|
| 274 | child.parent = self |
|---|
| 275 | |
|---|
| 276 | def insert(self, index, child): |
|---|
| 277 | """add a node to children""" |
|---|
| 278 | list_class.insert(self, index, child) |
|---|
| 279 | child.parent = self |
|---|
| 280 | |
|---|
| 281 | def remove(self, child): |
|---|
| 282 | """add a node to children""" |
|---|
| 283 | list_class.remove(self, child) |
|---|
| 284 | child.parent = None |
|---|
| 285 | |
|---|
| 286 | def pop(self, index): |
|---|
| 287 | """add a node to children""" |
|---|
| 288 | child = list_class.pop(self, index) |
|---|
| 289 | child.parent = None |
|---|
| 290 | |
|---|
| 291 | def __iter__(self): |
|---|
| 292 | return list_class.__iter__(self) |
|---|
| 293 | |
|---|
| 294 | # construct list from tree #################################################### |
|---|
| 295 | |
|---|
| 296 | def post_order_list(node, filter_func=no_filter): |
|---|
| 297 | """ |
|---|
| 298 | create a list with tree nodes for which the <filter> function returned true |
|---|
| 299 | in a post order fashion |
|---|
| 300 | """ |
|---|
| 301 | l, stack = [], [] |
|---|
| 302 | poped, index = 0, 0 |
|---|
| 303 | while node: |
|---|
| 304 | if filter_func(node): |
|---|
| 305 | if node.children and not poped: |
|---|
| 306 | stack.append((node, index)) |
|---|
| 307 | index = 0 |
|---|
| 308 | node = node.children[0] |
|---|
| 309 | else: |
|---|
| 310 | l.append(node) |
|---|
| 311 | index += 1 |
|---|
| 312 | try: |
|---|
| 313 | node = stack[-1][0].children[index] |
|---|
| 314 | except IndexError: |
|---|
| 315 | node = None |
|---|
| 316 | else: |
|---|
| 317 | node = None |
|---|
| 318 | poped = 0 |
|---|
| 319 | if node is None and stack: |
|---|
| 320 | node, index = stack.pop() |
|---|
| 321 | poped = 1 |
|---|
| 322 | return l |
|---|
| 323 | |
|---|
| 324 | def pre_order_list(node, filter_func=no_filter): |
|---|
| 325 | """ |
|---|
| 326 | create a list with tree nodes for which the <filter> function returned true |
|---|
| 327 | in a pre order fashion |
|---|
| 328 | """ |
|---|
| 329 | l, stack = [], [] |
|---|
| 330 | poped, index = 0, 0 |
|---|
| 331 | while node: |
|---|
| 332 | if filter_func(node): |
|---|
| 333 | if not poped: |
|---|
| 334 | l.append(node) |
|---|
| 335 | if node.children and not poped: |
|---|
| 336 | stack.append((node, index)) |
|---|
| 337 | index = 0 |
|---|
| 338 | node = node.children[0] |
|---|
| 339 | else: |
|---|
| 340 | index += 1 |
|---|
| 341 | try: |
|---|
| 342 | node = stack[-1][0].children[index] |
|---|
| 343 | except IndexError: |
|---|
| 344 | node = None |
|---|
| 345 | else: |
|---|
| 346 | node = None |
|---|
| 347 | poped = 0 |
|---|
| 348 | if node is None and len(stack) > 1: |
|---|
| 349 | node, index = stack.pop() |
|---|
| 350 | poped = 1 |
|---|
| 351 | return l |
|---|
| 352 | |
|---|
| 353 | class PostfixedDepthFirstIterator(FilteredIterator): |
|---|
| 354 | """a postfixed depth first iterator, designed to be used with visitors |
|---|
| 355 | """ |
|---|
| 356 | def __init__(self, node, filter_func=None): |
|---|
| 357 | FilteredIterator.__init__(self, node, post_order_list, filter_func) |
|---|
| 358 | |
|---|
| 359 | class PrefixedDepthFirstIterator(FilteredIterator): |
|---|
| 360 | """a pretfixed depth first iterator, designed to be used with visitors |
|---|
| 361 | """ |
|---|
| 362 | def __init__(self, node, filter_func=None): |
|---|
| 363 | FilteredIterator.__init__(self, node, pre_order_list, filter_func) |
|---|
| 364 | |
|---|