graph_drawing/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)