graph_drawing/line.py

def generate_line(root, node_type):
    """
    Print an expression tree in a single text line. Where needed, add
    parentheses.

    >>> from node import Node, Leaf
    >>> #from graph import generate_graph
    >>> l0, l1 = Leaf(1), Leaf(2)
    >>> n0 = Node('+', l0, l1)
    >>> print generate_line(n0, Node)
    1 + 2
    >>> n1 = Node('+', l0, l1)
    >>> n2 = Node('*', n0, n1)
    >>> print generate_line(n2, Node)
    (1 + 2) * (1 + 2)
    >>> l2 = Leaf(3)
    >>> n3 = Node('-', l2)
    >>> n4 = Node('*', n1, n3)
    >>> print generate_line(n4, Node)
    (1 + 2) * -3
    >>> #integral = Node('int', n0, n1)
    """

    operators = [
            ('+', '-'),
            ('*', '/'),
            ('^', )
            ]
    max_assoc = len(operators)

    def assoc(node):
        """
        Get the associativity of an operator node.
        """
        if isinstance(node, node_type) and len(node) > 1:
            op = node.title()

            for i, group in enumerate(operators):
                if op in group:
                    return i

        return max_assoc

    def traverse(node):
        """
        The expression tree is traversed using preorder traversal:
        1. Visit the root
        2. Traverse the left subtree
        3. Traverse the right subtree
        """
        s = node.title()

        if not isinstance(node, node_type):
            return s

        arity = len(node)

        if arity == 1:
            # Unary expression
            s += traverse(node[0])
        elif arity == 2:
            # Binary expression
            left, right = map(traverse, node)

            # Check if there is an assiociativity conflict on either side. If
            # so, add parentheses
            if assoc(node[0]) < assoc(node):
                left = '(' + left + ')'

            if assoc(node[1]) < assoc(node):
                right = '(' + right + ')'

            s = left + ' ' + s + ' ' + right

        return s

    return traverse(root)