math_parse.py

"""
My attempt at implementing a math parser, using the Shunting-yard algorithm
as described by its Wikipedia page:
https://en.wikipedia.org/wiki/Shunting-yard_algorithm
Currently it's a mess, and there are definitely things that I could simplify.
"""

import math

class Token:
    def __init__(self, value):
        self._val = value
    def __repr__(self):
        return str(self._val)
    def __str__(self):
        return str(self._val)
    def set_value(self, value):
        self._val = value
    def get_value(self):
        return self._val


class Value(Token):
    pass


class Var(Token):
    pass


class Op(Token):
    pass


class UnaryOp(Token):
    pass


class Coma(Token):
    pass


class Parenth(Token):
    pass


class Function(Token):
    pass


def get_variable(expr, i):
    var = ''
    while(i < len(expr) and any([
          expr[i].isnumeric(), expr[i].isalpha(),
          expr[i] == '_'])):
        var += expr[i]
        i += 1   
    return var, i


def get_value(expr, i):
    val = '' 
    while(i < len(expr) and 
         (expr[i].isnumeric()
         or expr[i] in ['.', 'e', 'E'])):
         if (expr[i] in ['e', 'E'] and 
             i+1 < len(expr) and 
             expr[i+1] in ['+', '-']):
             val += expr[i: i + 2]
             i += 2
         else:
             val += expr[i]
             i += 1
    val = float(val)
    return val, i


class Expression:

    functions = {'sin', 'cos', 'tan',
                 'sinh', 'cosh', 'tanh',
                 'arcsin', 'arccos', 'arctan',
                 'exp', 'log',
                 'max', 'min'}

    def __init__(self, expression):
        infix_tokens, variables = self._to_list(expression)
        self.infix_tokens = infix_tokens.copy()
        self._variables = variables
        self._rpn = self._infix_to_rpn(infix_tokens)

    @property
    def rpn(self):
        return self._rpn

    def __repr__(self):
        return str(self._rpn)

    def evalf(self, variables={}):
        func_key = {'sin': math.sin, 
                    'cos': math.cos, 
                    'tan': math.tan,
                    'sinh': math.sinh, 
                    'cosh': math.cosh, 
                    'tanh': math.tanh,
                    'arcsin': math.asin, 
                    'arccos': math.acos, 
                    'arctan': math.atan,
                    'exp': math.exp, 'log': math.log,
                    'max': max, 'min': min}
        for key in variables:
            if key in self._variables:
                self._variables[key].set_value(variables[key])
        eval_stack = []
        rpn = self._rpn.copy()
        while rpn:
            token = rpn.pop(0)
            if isinstance(token, Value) or isinstance(token, Var):
                eval_stack.append(token.get_value())
            elif isinstance(token, Op):
                val2 = eval_stack.pop()
                val1 = eval_stack.pop()
                val = 0
                if str(token) == '+':
                    val = val1 + val2
                elif str(token) == '-':
                    val = val1 - val2
                elif str(token) == '*':
                    val = val1*val2
                elif str(token) == '/':
                    val = val1/val2
                elif str(token) == '**' or str(token) == '^':
                    val = val1**val2
                eval_stack.append(val)
            elif isinstance(token, UnaryOp):
                val1 = eval_stack.pop()
                if str(token) == '-':
                    eval_stack.append(-val1)
            elif isinstance(token, Function):
                val1 = eval_stack.pop()
                # print(val1)    
                eval_stack.append(func_key[token.get_value()](val1))
            elif isinstance(token, Coma):
                val2 = eval_stack.pop()
                val1 = eval_stack.pop()
                list_val = []
                if isinstance(val2, list):
                    list_val.append(val1)
                    list_val.extend(val2)
                else:
                    list_val = [val1, val2]
                eval_stack.append(list_val)
            # print(eval_stack)
        val = eval_stack.pop()
        return val

    def _to_list(self, expr):
        tokens = []
        parenth_stack = []
        i = 0
        variables = {}
        while i < len(expr):
            if expr[i].isnumeric():
                val, i = get_value(expr, i)
                tokens.append(Value(val))
            elif expr[i].isalpha() or expr[i] == '_':
                var, i = get_variable(expr, i)
                if var in self.functions:
                    tokens.append(Function(var))
                else:
                    if not var in variables:
                        variables[var] = Var(var)
                    tokens.append(variables[var])
            else:
                if expr[i] in ['+', '-', '*', '/',
                            '^', '(', ')', ',']:
                    if (i + 1 < len(expr) and   
                        expr[i: i+2] == '**'):
                        tokens.append(Op('**'))
                        i += 2
                    elif expr[i] == ',':
                        tokens.append(Coma(','))
                        i += 1
                    else:
                        if expr[i] in ['(', ')']:
                            tokens.append(Parenth(expr[i]))
                            if expr[i] == '(':
                                parenth_stack.append('(')
                            else:
                                if parenth_stack:
                                    parenth_stack.pop()
                                else:
                                    raise SyntaxError('Unbalanced parenthesis')
                        else:
                            if not tokens or not (
                                    isinstance(tokens[-1], Var) or 
                                    isinstance(tokens[-1], Value) or
                                    str(tokens[-1]) == ')'):
                                # If statement is here so that a^-b is evaluated as
                                # a^(-b) while -a^b is evaluated as -(a^b).
                                # Not sure if entirely correct and there probably is a better
                                # way to do this.
                                if not tokens or (
                                    str(tokens[-1]) != '**' and str(tokens[-1]) != '^'
                                    ) and expr[i] == '-':
                                    tokens.append(Value(0.0))
                                    tokens.append(Op(expr[i]))
                                    # print(expr[i])
                                else:
                                    tokens.append(UnaryOp(expr[i]))
                            else:  
                                tokens.append(Op(expr[i]))
                        i += 1
                elif expr[i] == ' ':
                    i += 1
                else:
                    raise SyntaxError
                # print(tokens[-1], type(tokens[-1]))
        if parenth_stack:
            raise SyntaxError('Unbalanced parenthesis')
        return tokens, variables

    
    def _infix_to_rpn(self, expr_list):
        op_stack = []
        # parenth_stack = []
        rpn = []
        op_order = {'^': 0, '**': 0, '/': 1, '*': 2, '+': 3, '-': 3}
        while expr_list:
            item = expr_list.pop(0)
            if isinstance(item, Value) or isinstance(item, Var):
                rpn.append(item)
            elif isinstance(item, Function) or isinstance(item, UnaryOp):
                op_stack.append(item)
            elif isinstance(item, Op):
                while (op_stack and ((isinstance(op_stack[-1], Op)
                    and ((op_order[str(op_stack[-1])] != 0 and 
                    op_order[str(op_stack[-1])] <= op_order[str(item)]
                    ) or op_order[str(op_stack[-1])] < op_order[str(item)]))
                    or isinstance(op_stack[-1], Function)
                    or isinstance(op_stack[-1], UnaryOp))):
                    rpn.append(op_stack.pop())
                op_stack.append(item)
            elif isinstance(item, Coma):
                while str(op_stack[-1]) != ',' and str(op_stack[-1]) != '(':
                    rpn.append(op_stack.pop())
                op_stack.append(item)
            elif isinstance(item, Parenth) and str(item) == '(':
                op_stack.append(item)
            elif isinstance(item, Parenth) and str(item) == ')':
                while str(op_stack[-1]) != '(':
                    rpn.append(op_stack.pop())
                op_stack.pop()
        while op_stack:
            op = op_stack.pop()
            rpn.append(op)
        return rpn
 

print(expr := Expression('cos  (12.3434 + 11 - z) - 23e-1*11.0'))
print(expr.evalf({'z': math.pi}))
z = math.pi; print(math.cos  (12.3434 + 11 - z) - 23e-1*11.0)
print(expr := Expression('1.0-exp(-cos(-sin(pi)+3*tan(pi/4)))-23e-1*11.0'))
# print(expr.infix_tokens, [type(e) for e in expr.infix_tokens])
# print(expr.rpn, [type(e) for e in expr.rpn])
print(expr.evalf({'pi': math.pi}))
print(1.0-math.exp(-math.cos(-math.sin(math.pi) + 3*math.tan(math.pi/4)))-23e-1*11.0)
# print(expr := Expression('3 - 2 - 1'))
# print(expr.evalf())

# 1.0 pi sin - 3 pi 4 / tan * + cos - exp 2.3 11.0 * - -
# pi sin - 3 pi 4 / tan * + cos - exp - 23e-1 11.0 * -
# print(Expression('cos(z*y*t_12)**11 - a*b*sin(t)').rpn)
print(expr := Expression('cos(z*y*t_12)**11**12 - a*b*cosh(q)'))
z=12.0; y=0.1; t_12=34.34; a=0.1; b=2.0; q=3.14159; print(math.cos(z*y*t_12)**11**12 - a*b*math.cosh(q))
print(expr.evalf({'z': 12.0, 'y': 0.1, 't_12': 34.34, 'a': 0.1, 'b': 2.0, 'q': 3.14159}))
# z y t_12 * * cos 11 12 ** ** a b q cosh * * -
print(expr := Expression('a + b*z*11 - x*u/y'))
a=1.0; b=1.0/11.0; z=2.0; x=2.0; u=2.0; y=4.0; print(a + b*z*11 - x*u/y)
print(expr.evalf({'a': 1.0, 'b': 1.0/11.0, 'z': 2.0, 'x': 2.0, 'u': 2.0, 'y': 4.0}))
# a b z 11 * * + x u y / * -
print(expr := Expression('a*b/c - z*(a+b)'))
print(expr.evalf({'a':12.0, 'b': 0.5, 'c': 3, 'z': 1.0/12.5}))
a=12.0; b=0.5; c=3; z=1.0/12.5; print(a*b/c - z*(a+b))
# a b c / * z a b + * -
print(Expression('a*max(a, b+c*d, x - min(a, b, c))').rpn) 
# a a b c d * + x a b c , , min - , , max * 
# a ((a, (bcd*+), (x (a, b, c) min) - )) max * 
print(expr := Expression('max(a, b, c, d)')) 
print(expr.evalf({'a': -15.0, 'b': 10.0, 'c': -30.0, 'd': 6.0}))
expr = Expression('a + b*c^e')
print(expr.evalf({'a': 12.0, 'b': 0.1, 'c': 34.34, 'd': 0.1, 'e': 2.0}))
a=12.0; b=0.1; c=34.34; d=0.1; e=2.0; print(a+b*c**e)
print(expr.rpn)
expr = Expression('-a**b')
print(expr.rpn)