expr.py

"""
Simple infix mathematic expression calculator
Author: Julien Castiaux (github.com/Julien00859)
Licence: MIT
"""

from operator import add, sub, mul, truediv, mod

def shuntingyard(expr):
    """Convert infix expression to postfix"""
    ops = {"+": 0, "-": 0, "*": 1, "/": 1, "%": 1}
    op_stack = []
    ouput = []
    idx = 0
    length = len(expr)

    def number():
        nonlocal idx
        nbr = [char]
        while idx + 1 < length and (expr[idx + 1].isdigit() or expr[idx + 1] == "."):
            idx += 1
            nbr.append(expr[idx])
        ouput.append(float("".join(nbr)))

    def operator():
        while (op_stack
               and op_stack[-1] != "("
               and ops[op_stack[-1]] >= ops[char]):
            ouput.append(op_stack.pop())
        op_stack.append(char)

    def left_par():
        op_stack.append(char)

    def right_par():
        while op_stack and op_stack[-1] != "(":
            ouput.append(op_stack.pop())
        if not op_stack or op_stack.pop() != "(":
            raise SyntaxError("Unmatch parenthesis")

    while idx < length:
        char = expr[idx]
        if char in ops:
            operator()
        elif char.isdigit() or char == ".":
            number()
        elif char == "(":
            left_par()
        elif char == ")":
            right_par()
        elif char == " ":
            pass
        else:
            raise SyntaxError("Character '%s' is invalid" % char)
        idx += 1

    while op_stack:
        last_op = op_stack.pop()
        if last_op in "()":
            raise SyntaxError("Unmatch parenthesis")
        ouput.append(last_op)

    return ouput

def compute(expr):
    """Compute a postfix expression"""
    ops = {"+": add, "-": sub, "*": mul, "/": truediv, "%": mod}
    stack = []

    for element in expr:
        if isinstance(element, int, float):
            stack.append(element)
        else:
            if len(stack) < 2:
                raise SyntaxError("Missing operand")
            right = stack.pop()
            left = stack.pop()
            stack.append(ops[element](left, right))

    if len(stack) != 1:
        raise SyntaxError("Too many operators")
    return stack.pop()


if __name__ == "__main__":
    from sys import argv
    if "--test" in argv:
        assert compute(shuntingyard('15 / 3 + (1 + 3) * 2')) == 13
    else:
        infix_expr = " ".join(argv[1:])
        #print(infix_expr)
        postfix_expr = shuntingyard(infix_expr)
        #print(postfix_expr)
        result = compute(postfix_expr)
        print(result)