graph_drawing/line.py

from node import Leaf
from traverse import traverse_depth_first


OPERATORS = [
        ('+', '-'),
        ('*', '/', 'mod'),
        ('^', )
        ]

MAX_PRED = len(OPERATORS)


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

    return any(map(lambda x: label in x, OPERATORS))


def pred(node):
    """
    Get the precedence of an operator node.
    """
    # Check binary and n-ary operators
    if not node.is_leaf and len(node) > 1:
        op = node.title()

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

    # Unary operator and leaves have highest precedence
    return MAX_PRED


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

    >>> from node import Node, Leaf
    >>> l0, l1 = Leaf(1), Leaf(2)
    >>> plus = Node('+', l0, l1)
    >>> print generate_line_old(plus)
    1 + 2

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

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

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

    def traverse(node):
        """
        The expression tree is traversed using preorder traversal:
        1. Visit the root
        2. Traverse the subtrees in left-to-right order
        """
        if not node:
            return '<empty expression>'

        op = node.title()

        if not node.nodes:
            return op

        arity = len(node)

        if is_operator(node):
            if arity == 1:
                # Unary operator
                sub = node[0]
                sub_exp = traverse(sub)

                # Negated sub-expressions with spaces in them should be
                # enclosed in parentheses, unless they have a higher precedence
                # than subtraction are rewritten to a factor of a subtraction:
                # -(1 + 2)
                # -(1 - 2)
                # -4a
                # -(4 * 5)
                # 1 - 4 * 5
                # 1 + -(4 * 5)  ->  1 - 4 * 5
                if ' ' in sub_exp and sub.is_Leaf \
                        and hasattr(node, 'marked_negation') \
                        and pred(sub) > 0:
                    sub_exp = '(' + sub_exp + ')'

                result = op + sub_exp
            else:
                # N-ary operator
                node_pred = pred(node)
                result = ''
                sep = ' ' + op + ' '
                e = []

                # Mark added and subtracted negations for later use when adding
                # parentheses
                if op in ('+', '-'):
                    for child in node:
                        if child.title() == '-' and len(child) == 1:
                            child.marked_negation = True

                for i, child in enumerate(node):
                    exp = traverse(child)

                    # Check if there is a precedence conflict
                    # If so, add parentheses
                    child_pred = pred(child)

                    if child_pred < node_pred or \
                            (i and child_pred == node_pred \
                             and op != child.title()):
                        exp = '(' + exp + ')'

                    e.append(exp)

                if op == '*':
                    # Check if an explicit multiplication sign is nessecary
                    left, right = node

                    # Get the previous multiplication element if the arity is
                    # greater than 2
                    if left.title() == '*':
                        left = left[1]

                    # a * b -> ab
                    # a * 2 -> a * 2
                    # a * (b) -> a(b)
                    # (a) * b -> (a)b
                    # (a) * (b) -> (a)(b)
                    # 2 * a -> 2a
                    left_id = is_id(left)
                    right_id = is_id(right)
                    left_paren = e[0][-1] == ')'
                    right_paren = e[1][0] == '('
                    left_int = is_int(left)

                    if (left_id or left_paren or left_int) \
                            and (right_id or right_paren):
                        sep = ''

                result += sep.join(e)
        else:
            # Function call
            result = op + '(' + ', '.join(map(traverse, node)) + ')'

        return result

    # An addition with negation can be written as a subtraction, e.g.:
    # 1 + -2  ->  1 - 2
    return traverse(root).replace('+ -', '- ')


def is_id(node):
    return node.is_leaf and not node.title().isdigit()


def is_int(node):
    return node.is_leaf and node.title().isdigit()


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

    >>> from node import Node, 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)

    >>> minus = Node('-', Leaf(2), Node('-', Node('*', Leaf(15), Leaf('x'))))
    >>> print generate_line(minus)
    2 - -15x

    >>> left = Node('/', Leaf(22), Leaf(77))
    >>> right = Node('/', Leaf(28), Leaf(77))
    >>> minus = Node('-', left, right)
    >>> print generate_line(minus)
    22 / 77 - 28 / 77

    >>> plus = Node('+', left, Node('-', right))
    >>> print generate_line(plus)
    22 / 77 - 28 / 77
    """

    if not root:
        return '<empty expression>'

    if root.is_leaf:
        return '-' * root.negated + root.title()

    content = {}

    def construct_unary(node):
        result = node.title()
        value = node[0]

        # -a
        # -3*4
        # --a
        if value.is_leaf \
                or not ' ' in content[value] or pred(value) > 0:
            return result + content[value]

        return '%s(%s)' % (result, content[value])


    def construct_nary(node):
        op = node.title()

        # N-ary operator
        node_pred = pred(node)
        sep = ' ' + op + ' '
        e = []

        for i, child in enumerate(node):
            exp = content[child]

            # Check if there is a precedence conflict
            # If so, add parentheses
            child_pred = pred(child)

            if child_pred < node_pred or \
                    (i and child_pred == node_pred \
                        and op != child.title()):
                exp = '(' + exp + ')'

            e.append(exp)

        if op == '*':
            # Check if an explicit multiplication sign is nessecary
            left, right = node

            # Get the previous multiplication element if the arity is
            # greater than 2
            if left.title() == '*':
                left = left[1]

            # a * b -> ab
            # a * 2 -> a * 2
            # a * (b) -> a(b)
            # (a) * b -> (a)b
            # (a) * (b) -> (a)(b)
            # 2 * a -> 2a
            left_paren = e[0][-1] == ')'
            right_paren = e[1][0] == '('

            if (is_id(left) or left_paren or is_int(left)) \
                    and ((not right.negated and is_id(right)) or right_paren):
                sep = ''

        exp = sep.join(e)

        #if node.negated:
        # FIXME: Keep it this way?
        if node.negated and op != '*':
            exp = '(' + exp + ')'

        return exp

    def construct_function(node):
        buf = []

        for child in node:
            buf.append(content[child])

        return '%s(%s)' % (node.title(), ', '.join(buf))

    # Traverse the expression tree and construct the mathematical expression in
    # the leafs and nodes in depth first order.
    for node in traverse_depth_first(root):
        if node.is_leaf:
            content[node] = node.title()
        else:
            arity = len(node)

            if is_operator(node):
                if arity == 1:
                    content[node] = construct_unary(node)
                else:
                    content[node] = construct_nary(node)
            else:
                content[node] = construct_function(node)

        # Add negations
        content[node] = '-' * node.negated + content[node]

    # Merge binary plus and unary minus signs into binary minus.
    return content[root].replace('+ -', '- ')