Initial commit of text-based tree drawing.

.gitignore

+*.pyc
+*.sw[a-z]

graph.py

+# vim: set fileencoding=utf-8 :
+from node import Node, Leaf
+
+
+def generate_graph(root, node_type, separator=' '):
+    """
+    Return a text-based, utf-8 graph of a tree-like structure. Each tree node
+    is represented by a length-2 list. If a node has an attribute called
+    ``title``, that attribute will be called. That way, the node can return a
+    specific title, otherwise ``+`` is used.
+
+    >>> l0, l1 = Leaf(0), Leaf(1)
+    >>> n0 = Node('+', l0, l1)
+    >>> print generate_graph(n0, Node)
+     + 
+    ╭┴╮
+    0 1
+    >>> l2 = Leaf(2)
+    >>> n1 = Node('-', l2)
+    >>> print generate_graph(n1, Node)
+    -
+    │
+    2
+    >>> n2 = Node('*', n0, n1)
+    >>> print generate_graph(n2, Node)
+      +  
+     ╭┴─╮
+     +  -
+    ╭┴╮ │
+    0 1 2
+    """
+
+    node_width = {}
+
+    separator_len = len(separator)
+
+    def calculate_width(node):
+        title = node.title()
+        title_len = len(title)
+
+        # Leaves do not have children and therefore the length of its title is
+        # the width of the leaf.
+        if not isinstance(node, node_type):
+            node_width[node] = title_len
+            return title_len
+
+        width = 0
+
+        for child in node:
+            width += calculate_width(child)
+
+        # Add a separator between each node (thus n - 1 separators).
+        width += separator_len * len(node) - 1
+
+        # Odd numbers of children should be minus 1, since the middle child
+        # can be placed directly below the parent. With even numbers, there
+        # is no middle child, so the space below the parent cannot be used.
+        if len(node) % 2 == 1:
+            width -= 1
+
+        # If the title of the node is wider than the sum of its children, the
+        # title's width should be used.
+        width = max(title_len, width)
+
+        node_width[node] = width
+
+        return width
+
+    def node_middle(node):
+        node_len = len(node)
+        node_mid = node_len / 2
+
+        if node_len % 2 == 1:
+            node_mid += 1
+
+        middle = 0
+
+        for i, child in enumerate(node):
+            if i == node_mid:
+                break
+
+            middle += separator_len + node_width[child]
+
+        return middle - separator_len
+
+    def format_lines(node):
+        if not isinstance(node, node_type):
+            # Leaf titles do not need to be centered, since the parent will
+            # center those lines. And if there are no parents, the entire graph
+            # consists of a single leaf, so in that case there still is no
+            # reason to center it.
+            return [node.title()]
+
+        # At least one child, otherwise it would be a leaf.
+        assert node[0]
+
+        line_width = node_width[node]
+
+        lines = format_lines(node[0])
+
+        for child in node[1:]:
+            pos = len(lines[0])
+
+            child_lines = format_lines(child)
+            child_lines_len = len(child_lines)
+
+            # A node cannot have zero child lines.
+            assert child_lines
+
+            for i, line in enumerate(lines):
+                #print 'lines -> %d, "%s"' % (i, line)
+
+                if i < child_lines_len:
+                    padding_right = ' ' * (line_width - pos \
+                                           - len(child_lines[i]) \
+                                           - separator_len)
+
+                    lines[i] += separator + child_lines[i] + padding_right
+                else:
+                    # There are no more neighbor node on the right.
+                    lines[i] += ' '  * (line_width - pos)
+
+            # Add the child nodes that do not have neighbor nodes on the left.
+            for i, line in enumerate(child_lines[i+1:]):
+                #print 'child_lines[i:] -> %d, "%s"' % (i, line)
+                line = ' ' * (pos + separator_len) + line \
+                     + ' ' * (line_width - separator_len - pos - len(line))
+                lines.append(line)
+
+            # Validate that each line has an equal width
+            try:
+                for line in lines:
+                    assert len(line) == line_width
+            except AssertionError:
+
+                for l in lines:
+                    l = l.encode('utf-8')
+                    print l.replace(' ', '#'), len(l)
+
+                l = line.encode('utf-8')
+                print 'failed at:', l, len(l)
+                print 'current node:', node.title(), 'width:', line_width
+
+                raise
+
+        # Place the title above the middle two child nodes, or when there is a
+        # odd number of nodes, above the child node in the middle.
+        middle = node_middle(node)
+
+        #print 'node:', node.title(), 'middle:', middle, 'node_mid:', node_mid
+
+        title_line = center_text(node.title(), line_width, middle)
+
+        node_len = len(node)
+        node_mid = node_len / 2
+
+        edge_line = ''
+
+        for i, child in enumerate(node):
+            if isinstance(child, node_type):
+                middle = node_middle(child)
+            else:
+                middle = 0
+
+            if i < node_mid:
+                marker = ('╭' if i == 0 else '┬')
+                edge_line += center_text(marker.decode('utf-8'),
+                                        node_width[child],
+                                        middle, right='─'.decode('utf-8'))
+            else:
+                if i == node_mid:
+                    edge_line += '┴'.decode('utf-8')
+
+                marker = ('╮' if i == node_len - 1 else '┬')
+                edge_line += center_text(marker.decode('utf-8'),
+                                        node_width[child],
+                                        middle, left='─'.decode('utf-8'))
+
+        #try:
+        #    assert len(title_line) == len(edge_line)
+        #except AssertionError:
+        #    print 'title_line:', title_line, len(title_line)
+        #    print 'edge_line:', edge_line, len(edge_line)
+        #    raise
+
+        # Add the line of this node before all child lines.
+        return [title_line, edge_line] + lines
+
+    calculate_width(root)
+    lines = format_lines(root)
+
+    return '\n'.join(lines).encode('utf-8')
+
+def center_text(text, width, middle=0, left=' ', right=' '):
+    """
+    >>> center_text('+', 15, 11)
+    '           +   '
+    """
+    text_len = len(text)
+    text_mid = text_len / 2
+
+    # TODO: this code requires cleanup.
+
+    if middle:
+
+        # If this is true, the text is at the left.
+        if text_mid > middle:
+            text += left * (width - text_len)
+
+        # If this is true, the text is at the right.
+        elif text_mid > (width - middle):
+            text = right * (width - text_len) + text
+
+        # Else, the text has spacing on its left and right.
+        else:
+            text = left * (middle - text_mid) + text
+            text += right * (width - len(text))
+
+        return text
+
+    spacing = width - text_len
+
+    # Even number of spaces can be split equally.
+    if spacing % 2 == 0:
+        return left * (spacing / 2) + text + right * (spacing / 2)
+    # For an odd number of space, put the extra space at the end.
+    return left * (spacing / 2) + text + right * (spacing / 2 + 1)

node.py

+# vim: set fileencoding=utf-8 :
+
+class Node(object):
+    def __init__(self, label, *nodes):
+        self.label, self.nodes = label, nodes
+
+    def __getitem__(self, n):
+        return self.nodes[n]
+
+    def __setitem__(self, n, node):
+        self.nodes[n] = node
+
+    def __iter__(self):
+        return iter(self.nodes)
+
+    def __len__(self):
+        return len(self.nodes)
+
+    def title(self):
+        return str(self.label)
+
+
+class Leaf(object):
+    def __init__(self, label):
+        self.label = label
+
+    def title(self):
+        return str(self.label)