aboutsummaryrefslogtreecommitdiffstats
path: root/libbe/util/tree.py
diff options
context:
space:
mode:
Diffstat (limited to 'libbe/util/tree.py')
-rw-r--r--libbe/util/tree.py258
1 files changed, 258 insertions, 0 deletions
diff --git a/libbe/util/tree.py b/libbe/util/tree.py
new file mode 100644
index 0000000..812b0bd
--- /dev/null
+++ b/libbe/util/tree.py
@@ -0,0 +1,258 @@
+# Bugs Everywhere, a distributed bugtracker
+# Copyright (C) 2008-2010 Gianluca Montecchi <gian@grys.it>
+# W. Trevor King <wking@drexel.edu>
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; either version 2 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+"""Define :class:`Tree`, a traversable tree structure.
+"""
+
+import libbe
+if libbe.TESTING == True:
+ import doctest
+
+class Tree(list):
+ """A traversable tree structure.
+
+ Examples
+ --------
+
+ Construct::
+
+ +-b---d-g
+ a-+ +-e
+ +-c-+-f-h-i
+
+ with
+
+ >>> i = Tree(); i.n = "i"
+ >>> h = Tree([i]); h.n = "h"
+ >>> f = Tree([h]); f.n = "f"
+ >>> e = Tree(); e.n = "e"
+ >>> c = Tree([f,e]); c.n = "c"
+ >>> g = Tree(); g.n = "g"
+ >>> d = Tree([g]); d.n = "d"
+ >>> b = Tree([d]); b.n = "b"
+ >>> a = Tree(); a.n = "a"
+ >>> a.append(c)
+ >>> a.append(b)
+
+ Get the longest branch length with
+
+ >>> a.branch_len()
+ 5
+
+ Sort the tree recursively. Here we sort longest branch length
+ first.
+
+ >>> a.sort(key=lambda node : -node.branch_len())
+ >>> "".join([node.n for node in a.traverse()])
+ 'acfhiebdg'
+
+ And here we sort shortest branch length first.
+
+ >>> a.sort(key=lambda node : node.branch_len())
+ >>> "".join([node.n for node in a.traverse()])
+ 'abdgcefhi'
+
+ We can also do breadth-first traverses.
+
+ >>> "".join([node.n for node in a.traverse(depth_first=False)])
+ 'abcdefghi'
+
+ Serialize the tree with depth marking branches.
+
+ >>> for depth,node in a.thread():
+ ... print "%*s" % (2*depth+1, node.n)
+ a
+ b
+ d
+ g
+ c
+ e
+ f
+ h
+ i
+
+ Flattening the thread disables depth increases except at
+ branch splits.
+
+ >>> for depth,node in a.thread(flatten=True):
+ ... print "%*s" % (2*depth+1, node.n)
+ a
+ b
+ d
+ g
+ c
+ e
+ f
+ h
+ i
+
+ We can also check if a node is contained in a tree.
+
+ >>> a.has_descendant(g)
+ True
+ >>> c.has_descendant(g)
+ False
+ >>> a.has_descendant(a)
+ False
+ >>> a.has_descendant(a, match_self=True)
+ True
+ """
+ def __cmp__(self, other):
+ return cmp(id(self), id(other))
+
+ def __eq__(self, other):
+ return self.__cmp__(other) == 0
+
+ def __ne__(self, other):
+ return self.__cmp__(other) != 0
+
+ def branch_len(self):
+ """Return the largest number of nodes from root to leaf (inclusive).
+
+ For the tree::
+
+ +-b---d-g
+ a-+ +-e
+ +-c-+-f-h-i
+
+ this method returns 5.
+
+ Notes
+ -----
+ Exhaustive search every time == *slow*.
+
+ Use only on small trees, or reimplement by overriding
+ child-addition methods to allow accurate caching.
+ """
+ if len(self) == 0:
+ return 1
+ else:
+ return 1 + max([child.branch_len() for child in self])
+
+ def sort(self, *args, **kwargs):
+ """Sort the tree recursively.
+
+ This method extends :meth:`list.sort` to Trees.
+
+ Notes
+ -----
+ This method can be slow, e.g. on a :meth:`branch_len` sort,
+ since a node at depth `N` from the root has it's
+ :meth:`branch_len` method called `N` times.
+ """
+ list.sort(self, *args, **kwargs)
+ for child in self:
+ child.sort(*args, **kwargs)
+
+ def traverse(self, depth_first=True):
+ """Generate all the nodes in a tree, starting with the root node.
+
+ Parameters
+ ----------
+ depth_first : bool
+ Depth first by default, but you can set `depth_first` to
+ `False` for breadth first ordering. Siblings are returned
+ in the order they are stored, so you might want to
+ :meth:`sort` your tree first.
+ """
+ if depth_first == True:
+ yield self
+ for child in self:
+ for descendant in child.traverse():
+ yield descendant
+ else: # breadth first, Wikipedia algorithm
+ # http://en.wikipedia.org/wiki/Breadth-first_search
+ queue = [self]
+ while len(queue) > 0:
+ node = queue.pop(0)
+ yield node
+ queue.extend(node)
+
+ def thread(self, flatten=False):
+ """Generate a (depth, node) tuple for every node in the tree.
+
+ When `flatten` is `False`, the depth of any node is one
+ greater than the depth of its parent. That way the
+ inheritance is explicit, but you can end up with highly
+ indented threads.
+
+ When `flatten` is `True`, the depth of any node is only
+ greater than the depth of its parent when there is a branch,
+ and the node is not the last child. This can lead to ancestry
+ ambiguity, but keeps the total indentation down. For example::
+
+ +-b +-b-c
+ a-+-c and a-+
+ +-d-e-f +-d-e-f
+
+ would both produce (after sorting by :meth:`branch_len`)::
+
+ (0, a)
+ (1, b)
+ (1, c)
+ (0, d)
+ (0, e)
+ (0, f)
+
+ """
+ stack = [] # ancestry of the current node
+ if flatten == True:
+ depthDict = {}
+
+ for node in self.traverse(depth_first=True):
+ while len(stack) > 0 \
+ and id(node) not in [id(c) for c in stack[-1]]:
+ stack.pop(-1)
+ if flatten == False:
+ depth = len(stack)
+ else:
+ if len(stack) == 0:
+ depth = 0
+ else:
+ parent = stack[-1]
+ depth = depthDict[id(parent)]
+ if len(parent) > 1 and node != parent[-1]:
+ depth += 1
+ depthDict[id(node)] = depth
+ yield (depth,node)
+ stack.append(node)
+
+ def has_descendant(self, descendant, depth_first=True, match_self=False):
+ """Check if a node is contained in a tree.
+
+ Parameters
+ ----------
+ descendant : Tree
+ The potential descendant.
+ depth_first : bool
+ The search order. Set this if you feel depth/breadth would
+ be a faster search.
+ match_self : bool
+ Set to `True` for::
+
+ x.has_descendant(x, match_self=True) -> True
+ """
+ if descendant == self:
+ return match_self
+ for d in self.traverse(depth_first):
+ if descendant == d:
+ return True
+ return False
+
+if libbe.TESTING == True:
+ suite = doctest.DocTestSuite()