graph_drawing/line.py

from node import Node

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

    >>> from node import Leaf

    >>> l0, l1 = Leaf(1), Leaf(2)
    >>> plus = Node('+', l0, l1)
    >>> print generate_line(plus)
    1 + 2

    >>> plus2 = Node('+', l0, l1)
    >>> times = Node('*', plus, plus2)
    >>> print generate_line(times)
    (1 + 2) * (1 + 2)

    >>> l2 = Leaf(3)
    >>> uminus = Node('-', l2)
    >>> times = Node('*', plus, uminus)
    >>> print generate_line(times)
    (1 + 2) * -3

    >>> exp = Leaf('x')
    >>> inf = Leaf('oo')
    >>> minus_inf = Node('-', inf)
    >>> integral = Node('int', exp, minus_inf, inf)
    >>> print generate_line(integral)
    int(x, -oo, oo)
    """

    # FIXME: Binding order
    operators = [
            ('+', '-'),
            ('mod', ),
            ('*', '/'),
            ('^', )
            ]
    max_assoc = len(operators)

    def is_operator(node):
        """
        Check if a given node is an operator (otherwise, it's a function).
        """
        label = node.title()
        either = lambda a, b: a or b

        return reduce(either, map(lambda x: label in x, operators))

    def assoc(node):
        """
        Get the associativity of an operator node.
        """
        if isinstance(node, Node) 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 subtrees in left-to-right order
        """
        s = node.title()

        if not isinstance(node, Node):
            return s

        arity = len(node)

        if is_operator(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
                node_assoc = assoc(node)

                if assoc(node[0]) < node_assoc:
                    left = '(' + left + ')'

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

                s = left + ' ' + s + ' ' + right
            else:  # pragma: nocover
                raise ValueError('arity = %d is currently not supported.' \
                                 % arity)
        else:
            # Function
            s += '(' + ', '.join(map(traverse, node)) + ')'

        return s

    return traverse(root)