graph_drawing/graph.py

# vim: set fileencoding=utf-8 :
# XXX Used in doctests (we should use them in the __main__ section below too).
def generate_graph(root, separator=' ', verbose=False):
    """
    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.

    >>> from node import Leaf, Node
    >>> l0, l1 = Leaf(0), Leaf(1)
    >>> n0 = Node('+', l0, l1)
    >>> l2 = Leaf(2)
    >>> print generate_graph(n0)
     +
    ╭┴╮
    0 1
    >>> n1 = Node('-', l2)
    >>> print generate_graph(n1)
    -
    │
    2
    >>> n2 = Node('*', n0, n1)
    >>> print generate_graph(n2)
       *
     ╭─┴╮
     +  -
    ╭┴╮ │
    0 1 2
    """

    node_width = {}
    node_middle = {}

    separator_len = len(separator)

    def calculate_node_sizes(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 node.nodes:
            node_width[node] = title_len
            node_middle[node] = int((title_len - 1) / 2)
            return title_len

        node_len = len(node)

        width = 0
        middle = 0
        middle_pos = int(node_len / 2)

        for i, child in enumerate(node):
            tmp = calculate_node_sizes(child)

            if i < middle_pos:
                middle += tmp

            width += tmp

        middle += max(middle_pos - int(node_len % 2 == 0), 0) * separator_len

        # Add a separator between each node (thus n - 1 separators).
        width += separator_len * (node_len - 1)

        # If the title of the node is wider than the sum of its children, the
        # title's width should be used.
        if title_len > width:
            width = title_len
            middle = int(title_len / 2)

        # print 'width of "%s":' % node.title(), width

        node_width[node] = width
        node_middle[node] = middle

        return width

    def format_lines(node):
        if not node.nodes:
            # 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]

        child_lines = [format_lines(child) for child in node]
        max_height = max(map(len, child_lines))

        # Assert that all child boxes are of equal height
        for lines in child_lines:
            additional_line = separator * len(lines[0])
            lines += [additional_line for i in range(max_height - len(lines))]

        assert len(child_lines[0]) == max_height
        from copy import deepcopy
        result = deepcopy(child_lines[0])

        for lines in child_lines[1:]:
            assert len(lines) == max_height

            for i, line in enumerate(lines):
                result[i] += separator + line

        line_width = node_width[node]

        # TODO: substitute box_widths with node_width
        box_widths = [len(lines[0]) for lines in child_lines]
        node_len = len(node)
        middle_node = int(node_len / 2)
        #middle = sum([box_widths[i] for i in range(middle_node)]) \
        #         + max(middle_node - int(node_len % 2 == 0), 0) * separator_len
        middle = node_middle[node]

        title_line = center_text(node.title(), line_width, middle)

        pipe_sign = '│'.decode('utf-8')
        dash_sign = '─'.decode('utf-8')
        cross_sign = '┼'.decode('utf-8')
        tsplit_dn_sign = '┬'.decode('utf-8')
        tsplit_up_sign = '┴'.decode('utf-8')
        left_sign = '╭'.decode('utf-8')
        right_sign = '╮'.decode('utf-8')

        if node_len == 1:
            # Unary operators
            edge_line = center_text(pipe_sign, box_widths[0], middle)
        elif node_len % 2:
            # n-ary operators (n is odd)
            edge_line = ''

            for i, child in enumerate(node):
                if i > 0:
                    edge_line += dash_sign

                if i < middle_node:
                    marker = (left_sign if i == 0 else tsplit_dn_sign)
                    edge_line += center_text(marker, box_widths[i],
                                             middle=0, right=dash_sign)
                else:
                    if i == middle_node:
                        marker = cross_sign
                        edge_line += center_text(marker, box_widths[i],
                                                middle=0, right=dash_sign,
                                                left=dash_sign)
                    else:
                        if i == node_len - 1:
                            marker = right_sign
                        else:
                            marker = tsplit_dn_sign

                        edge_line += center_text(marker, box_widths[i],
                                                 middle=0, left=dash_sign)
        else:
            # n-ary operators (n is even)
            edge_line = ''

            for i, child in enumerate(node):
                if i > 0:
                    if i == middle_node:
                        edge_line += tsplit_up_sign
                    else:
                        edge_line += dash_sign

                if i < middle_node:
                    marker = (left_sign if i == 0 else tsplit_dn_sign)
                    edge_line += center_text(marker, box_widths[i],
                                             middle=node_middle[child],
                                             right=dash_sign)
                else:
                    if i == node_len - 1:
                        marker = right_sign
                    else:
                        marker = tsplit_dn_sign

                    edge_line += center_text(marker, box_widths[i],
                                             middle=node_middle[child],
                                             left=dash_sign)

        try:
            assert len(title_line) == len(edge_line)
        except AssertionError:  # pragma: nocover
            print '------------------'
            print 'line_width:', line_width
            print 'title_line:', title_line, 'len:', len(title_line)
            print 'edge_line: %s (%d)' % (edge_line.encode('utf-8'),
                                          len(edge_line))
            print 'lines:'
            print '\n'.join(map(lambda x: x + ' %d' % len(x), lines))
            raise Exception()

        # Add the line of this node before all child lines.
        return [title_line, edge_line] + result

    calculate_node_sizes(root)

    if verbose:  # pragma: nocover
        print '------- node_{width,middle} ---------'
        for node, width in node_width.iteritems():
            print node.title(), 'width:', width, 'middle:', node_middle[node]

    lines = format_lines(root)

    # Strip trailing whitespace.
    return '\n'.join(map(lambda x: x.rstrip(), lines)).encode('utf-8')


def center_text(text, width, middle=0, left=' ', right=' '):
    """
    >>> print center_text('│', 1, 1)
    │
    >>> center_text('+', 15, 11)
    '           +   '
    >>> left = center_text('╭'.decode('utf-8'), 11, 8, right='─'.decode('utf-8'))
    >>> len(left) == 11
    True
    >>> right = center_text('╮'.decode('utf-8'), 3, 2, left='─'.decode('utf-8'))
    >>> len(right) == 3
    True
    >>> edge_line = left + '┴'.decode('utf-8') + right
    >>> len(edge_line) == 15
    True
    >>> title_line = center_text('+', 15, 11)
    >>> print '|%s|\\n|%s|' % (title_line, edge_line.encode('utf-8'))
    |           +   |
    |        ╭──┴──╮|
    """
    text_len = len(text)
    text_mid = text_len / 2

    #print '---------'
    #print 'text_len:', text_len
    #print 'text_mid:', text_mid
    #print 'width:', width
    #print 'middle:', middle
    #print '---------'

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