Merged Trevor's tree

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()