libbe/util/tree.py

   1 # Bugs Everywhere, a distributed bugtracker
   2 # Copyright (C) 2008-2011 Chris Ball <cjb@laptop.org>
   3 #                         Gianluca Montecchi <gian@grys.it>
   4 #                         W. Trevor King <wking@drexel.edu>
   5 #
   6 # This file is part of Bugs Everywhere.
   7 #
   8 # Bugs Everywhere is free software; you can redistribute it and/or modify it
   9 # under the terms of the GNU General Public License as published by the
  10 # Free Software Foundation, either version 2 of the License, or (at your
  11 # option) any later version.
  12 #
  13 # Bugs Everywhere is distributed in the hope that it will be useful, but
  14 # WITHOUT ANY WARRANTY; without even the implied warranty of
  15 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  16 # General Public License for more details.
  17 #
  18 # You should have received a copy of the GNU General Public License
  19 # along with Bugs Everywhere.  If not, see <http://www.gnu.org/licenses/>.
  20
  21 """Define :class:`Tree`, a traversable tree structure.
  22 """
  23
  24 import libbe
  25 if libbe.TESTING == True:
  26     import doctest
  27
  28 class Tree(list):
  29     """A traversable tree structure.
  30
  31     Examples
  32     --------
  33
  34     Construct::
  35
  36                +-b---d-g
  37              a-+   +-e
  38                +-c-+-f-h-i
  39
  40     with
  41
  42     >>> i = Tree();       i.n = "i"
  43     >>> h = Tree([i]);    h.n = "h"
  44     >>> f = Tree([h]);    f.n = "f"
  45     >>> e = Tree();       e.n = "e"
  46     >>> c = Tree([f,e]);  c.n = "c"
  47     >>> g = Tree();       g.n = "g"
  48     >>> d = Tree([g]);    d.n = "d"
  49     >>> b = Tree([d]);    b.n = "b"
  50     >>> a = Tree();       a.n = "a"
  51     >>> a.append(c)
  52     >>> a.append(b)
  53
  54     Get the longest branch length with
  55
  56     >>> a.branch_len()
  57     5
  58
  59     Sort the tree recursively.  Here we sort longest branch length
  60     first.
  61
  62     >>> a.sort(key=lambda node : -node.branch_len())
  63     >>> "".join([node.n for node in a.traverse()])
  64     'acfhiebdg'
  65
  66     And here we sort shortest branch length first.
  67
  68     >>> a.sort(key=lambda node : node.branch_len())
  69     >>> "".join([node.n for node in a.traverse()])
  70     'abdgcefhi'
  71
  72     We can also do breadth-first traverses.
  73
  74     >>> "".join([node.n for node in a.traverse(depth_first=False)])
  75     'abcdefghi'
  76
  77     Serialize the tree with depth marking branches.
  78
  79     >>> for depth,node in a.thread():
  80     ...     print "%*s" % (2*depth+1, node.n)
  81     a
  82       b
  83         d
  84           g
  85       c
  86         e
  87         f
  88           h
  89             i
  90
  91     Flattening the thread disables depth increases except at
  92     branch splits.
  93
  94     >>> for depth,node in a.thread(flatten=True):
  95     ...     print "%*s" % (2*depth+1, node.n)
  96     a
  97       b
  98       d
  99       g
 100     c
 101       e
 102     f
 103     h
 104     i
 105
 106     We can also check if a node is contained in a tree.
 107
 108     >>> a.has_descendant(g)
 109     True
 110     >>> c.has_descendant(g)
 111     False
 112     >>> a.has_descendant(a)
 113     False
 114     >>> a.has_descendant(a, match_self=True)
 115     True
 116     """
 117     def __cmp__(self, other):
 118         return cmp(id(self), id(other))
 119
 120     def __eq__(self, other):
 121         return self.__cmp__(other) == 0
 122
 123     def __ne__(self, other):
 124         return self.__cmp__(other) != 0
 125
 126     def branch_len(self):
 127         """Return the largest number of nodes from root to leaf (inclusive).
 128
 129         For the tree::
 130
 131                +-b---d-g
 132              a-+   +-e
 133                +-c-+-f-h-i
 134
 135         this method returns 5.
 136
 137         Notes
 138         -----
 139         Exhaustive search every time == *slow*.
 140
 141         Use only on small trees, or reimplement by overriding
 142         child-addition methods to allow accurate caching.
 143         """
 144         if len(self) == 0:
 145             return 1
 146         else:
 147             return 1 + max([child.branch_len() for child in self])
 148
 149     def sort(self, *args, **kwargs):
 150         """Sort the tree recursively.
 151
 152         This method extends :meth:`list.sort` to Trees.
 153
 154         Notes
 155         -----
 156         This method can be slow, e.g. on a :meth:`branch_len` sort,
 157         since a node at depth `N` from the root has it's
 158         :meth:`branch_len` method called `N` times.
 159         """
 160         list.sort(self, *args, **kwargs)
 161         for child in self:
 162             child.sort(*args, **kwargs)
 163
 164     def traverse(self, depth_first=True):
 165         """Generate all the nodes in a tree, starting with the root node.
 166
 167         Parameters
 168         ----------
 169         depth_first : bool
 170           Depth first by default, but you can set `depth_first` to
 171           `False` for breadth first ordering.  Siblings are returned
 172           in the order they are stored, so you might want to
 173           :meth:`sort` your tree first.
 174         """
 175         if depth_first == True:
 176             yield self
 177             for child in self:
 178                 for descendant in child.traverse():
 179                     yield descendant
 180         else: # breadth first, Wikipedia algorithm
 181             # http://en.wikipedia.org/wiki/Breadth-first_search
 182             queue = [self]
 183             while len(queue) > 0:
 184                 node = queue.pop(0)
 185                 yield node
 186                 queue.extend(node)
 187
 188     def thread(self, flatten=False):
 189         """Generate a (depth, node) tuple for every node in the tree.
 190
 191         When `flatten` is `False`, the depth of any node is one
 192         greater than the depth of its parent.  That way the
 193         inheritance is explicit, but you can end up with highly
 194         indented threads.
 195
 196         When `flatten` is `True`, the depth of any node is only
 197         greater than the depth of its parent when there is a branch,
 198         and the node is not the last child.  This can lead to ancestry
 199         ambiguity, but keeps the total indentation down.  For example::
 200
 201                       +-b                  +-b-c
 202                     a-+-c        and     a-+
 203                       +-d-e-f              +-d-e-f
 204
 205         would both produce (after sorting by :meth:`branch_len`)::
 206
 207             (0, a)
 208             (1, b)
 209             (1, c)
 210             (0, d)
 211             (0, e)
 212             (0, f)
 213
 214         """
 215         stack = [] # ancestry of the current node
 216         if flatten == True:
 217             depthDict = {}
 218
 219         for node in self.traverse(depth_first=True):
 220             while len(stack) > 0 \
 221                     and id(node) not in [id(c) for c in stack[-1]]:
 222                 stack.pop(-1)
 223             if flatten == False:
 224                 depth = len(stack)
 225             else:
 226                 if len(stack) == 0:
 227                     depth = 0
 228                 else:
 229                     parent = stack[-1]
 230                     depth = depthDict[id(parent)]
 231                     if len(parent) > 1 and node != parent[-1]:
 232                         depth += 1
 233                 depthDict[id(node)] = depth
 234             yield (depth,node)
 235             stack.append(node)
 236
 237     def has_descendant(self, descendant, depth_first=True, match_self=False):
 238         """Check if a node is contained in a tree.
 239
 240         Parameters
 241         ----------
 242         descendant : Tree
 243           The potential descendant.
 244         depth_first : bool
 245           The search order.  Set this if you feel depth/breadth would
 246           be a faster search.
 247         match_self : bool
 248           Set to `True` for::
 249
 250               x.has_descendant(x, match_self=True) -> True
 251         """
 252         if descendant == self:
 253             return match_self
 254         for d in self.traverse(depth_first):
 255             if descendant == d:
 256                 return True
 257         return False
 258
 259 if libbe.TESTING == True:
 260     suite = doctest.DocTestSuite()