libbe/util/tree.py

   1 # Bugs Everywhere, a distributed bugtracker
   2 # Copyright (C) 2008-2009 Gianluca Montecchi <gian@grys.it>
   3 #                         W. Trevor King <wking@drexel.edu>
   4 #
   5 # This program is free software; you can redistribute it and/or modify
   6 # it under the terms of the GNU General Public License as published by
   7 # the Free Software Foundation; either version 2 of the License, or
   8 # (at your option) any later version.
   9 #
  10 # This program is distributed in the hope that it will be useful,
  11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 # GNU General Public License for more details.
  14 #
  15 # You should have received a copy of the GNU General Public License along
  16 # with this program; if not, write to the Free Software Foundation, Inc.,
  17 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  18
  19 """
  20 Define a traversable tree structure.
  21 """
  22
  23 import libbe
  24 if libbe.TESTING == True:
  25     import doctest
  26
  27 class Tree(list):
  28     """
  29     Construct
  30                +-b---d-g
  31              a-+   +-e
  32                +-c-+-f-h-i
  33     with
  34     >>> i = Tree();       i.n = "i"
  35     >>> h = Tree([i]);    h.n = "h"
  36     >>> f = Tree([h]);    f.n = "f"
  37     >>> e = Tree();       e.n = "e"
  38     >>> c = Tree([f,e]);  c.n = "c"
  39     >>> g = Tree();       g.n = "g"
  40     >>> d = Tree([g]);    d.n = "d"
  41     >>> b = Tree([d]);    b.n = "b"
  42     >>> a = Tree();       a.n = "a"
  43     >>> a.append(c)
  44     >>> a.append(b)
  45
  46     >>> a.branch_len()
  47     5
  48     >>> a.sort(key=lambda node : -node.branch_len())
  49     >>> "".join([node.n for node in a.traverse()])
  50     'acfhiebdg'
  51     >>> a.sort(key=lambda node : node.branch_len())
  52     >>> "".join([node.n for node in a.traverse()])
  53     'abdgcefhi'
  54     >>> "".join([node.n for node in a.traverse(depth_first=False)])
  55     'abcdefghi'
  56     >>> for depth,node in a.thread():
  57     ...     print "%*s" % (2*depth+1, node.n)
  58     a
  59       b
  60         d
  61           g
  62       c
  63         e
  64         f
  65           h
  66             i
  67     >>> for depth,node in a.thread(flatten=True):
  68     ...     print "%*s" % (2*depth+1, node.n)
  69     a
  70       b
  71       d
  72       g
  73     c
  74       e
  75     f
  76     h
  77     i
  78     >>> a.has_descendant(g)
  79     True
  80     >>> c.has_descendant(g)
  81     False
  82     >>> a.has_descendant(a)
  83     False
  84     >>> a.has_descendant(a, match_self=True)
  85     True
  86     """
  87     def __cmp__(self, other):
  88         return cmp(id(self), id(other))
  89
  90     def __eq__(self, other):
  91         return self.__cmp__(other) == 0
  92
  93     def __ne__(self, other):
  94         return self.__cmp__(other) != 0
  95
  96     def branch_len(self):
  97         """
  98         Exhaustive search every time == SLOW.
  99
 100         Use only on small trees, or reimplement by overriding
 101         child-addition methods to allow accurate caching.
 102
 103         For the tree
 104                +-b---d-g
 105              a-+   +-e
 106                +-c-+-f-h-i
 107         this method returns 5.
 108         """
 109         if len(self) == 0:
 110             return 1
 111         else:
 112             return 1 + max([child.branch_len() for child in self])
 113
 114     def sort(self, *args, **kwargs):
 115         """
 116         This method can be slow, e.g. on a branch_len() sort, since a
 117         node at depth N from the root has it's branch_len() method
 118         called N times.
 119         """
 120         list.sort(self, *args, **kwargs)
 121         for child in self:
 122             child.sort(*args, **kwargs)
 123
 124     def traverse(self, depth_first=True):
 125         """
 126         Note: you might want to sort() your tree first.
 127         """
 128         if depth_first == True:
 129             yield self
 130             for child in self:
 131                 for descendant in child.traverse():
 132                     yield descendant
 133         else: # breadth first, Wikipedia algorithm
 134             # http://en.wikipedia.org/wiki/Breadth-first_search
 135             queue = [self]
 136             while len(queue) > 0:
 137                 node = queue.pop(0)
 138                 yield node
 139                 queue.extend(node)
 140
 141     def thread(self, flatten=False):
 142         """
 143         When flatten==False, the depth of any node is one greater than
 144         the depth of its parent.  That way the inheritance is
 145         explicit, but you can end up with highly indented threads.
 146
 147         When flatten==True, the depth of any node is only greater than
 148         the depth of its parent when there is a branch, and the node
 149         is not the last child.  This can lead to ancestry ambiguity,
 150         but keeps the total indentation down.  E.g.
 151                       +-b                  +-b-c
 152                     a-+-c        and     a-+
 153                       +-d-e-f              +-d-e-f
 154         would both produce (after sorting by branch_len())
 155         (0, a)
 156         (1, b)
 157         (1, c)
 158         (0, d)
 159         (0, e)
 160         (0, f)
 161         """
 162         stack = [] # ancestry of the current node
 163         if flatten == True:
 164             depthDict = {}
 165
 166         for node in self.traverse(depth_first=True):
 167             while len(stack) > 0 \
 168                     and id(node) not in [id(c) for c in stack[-1]]:
 169                 stack.pop(-1)
 170             if flatten == False:
 171                 depth = len(stack)
 172             else:
 173                 if len(stack) == 0:
 174                     depth = 0
 175                 else:
 176                     parent = stack[-1]
 177                     depth = depthDict[id(parent)]
 178                     if len(parent) > 1 and node != parent[-1]:
 179                         depth += 1
 180                 depthDict[id(node)] = depth
 181             yield (depth,node)
 182             stack.append(node)
 183
 184     def has_descendant(self, descendant, depth_first=True, match_self=False):
 185         if descendant == self:
 186             return match_self
 187         for d in self.traverse(depth_first):
 188             if descendant == d:
 189                 return True
 190         return False
 191
 192 if libbe.TESTING == True:
 193     suite = doctest.DocTestSuite()