Arithmetic Expression Evaluation

.gitignore

*.md

eval.py

# This is to implement the idea I posted in my [StackOverflow answer](https://stackoverflow.com/questions/69349502/simple-program-to-perform-a-math-formula-in-python/69350338#69350338).
#
# Sample input:
# (a+b)/c
# a=2;b=56;c=5
#
# We can assume that all variables are single letter, and there are only four operators: +-*/
#
# Expected output:
# (a+b)/c
# (2+56)/5
# 58/5
# 11.6

from enum import Enum, auto

class TokenType(Enum):
    VARIABLE = auto()
    OPERATOR = auto()
    LEFT_PARENTHESIS = auto()
    RIGHT_PARENTHESIS = auto()

class Token:
    # for variable and operator, `val` is the character
    # for parentheses, `val` is the index of the other half
    def __init__(self, type, value):
        self.type = type
        self.value = value
    def __repr__(self):
        return self.__str__()
    def __str__(self):
        return f"{self.type}: {self.value}"

class NodeType(Enum):
    NUMBER = auto()
    VARIABLE = auto()
    OPERATOR = auto()

class SyntaxNode:
    # `parentheses_wrapped` is for outputing parentheses only
    # if we only care about the final result then it's not necessary
    def __init__(self, type, value, left=None, right=None, parentheses_wrapped=False):
        self.type = type
        self.value = value
        self.left = left
        self.right = right
        self.parentheses_wrapped = parentheses_wrapped

    def evaluate(self, context, callback=None):
        if self.type == NodeType.NUMBER:
            return self.value # no need to call callback here
        elif self.type == NodeType.VARIABLE:
            # replace with number node
            if not self.value in context:
                raise ValueError(f"Can't find value for variable {self.value}")
            num = context[self.value]
            self.type = NodeType.NUMBER
            self.value = num
            if callback:
                callback()
            return self.value
        elif self.type == NodeType.OPERATOR:
            new_val = None
            if self.value == '+':
                new_val = self.left.evaluate(context, callback) + self.right.evaluate(context, callback)
            elif self.value == '-':
                new_val = self.left.evaluate(context, callback) - self.right.evaluate(context, callback)
            elif self.value == '*':
                new_val = self.left.evaluate(context, callback) * self.right.evaluate(context, callback)
            elif self.value == '/':
                new_val = self.left.evaluate(context, callback) / self.right.evaluate(context, callback)
            else:
                raise NameError('Unknown operator: {}'.format(self.val))
            self.type = NodeType.NUMBER
            self.value = new_val
            self.left = None
            self.right = None
            if self.value >= 0:
                self.parentheses_wrapped = False
            if callback:
                callback()
            return new_val

def parse_expression_to_tokens(exp):
    n = len(exp)
    tokens = []
    stack = []
    for i in range(n):
        c = exp[i]
        if c == '(':
            token = Token(TokenType.LEFT_PARENTHESIS, i) # save own index for now. change later
            stack.append(token)
            tokens.append(token)
        elif c == ')':
            token = Token(TokenType.RIGHT_PARENTHESIS, -1)
            l_token = stack.pop()
            token.value = l_token.value
            l_token.value = i
            tokens.append(token)
        elif c == '+' or c == '-' or c == '*' or c == '/':
            tokens.append(Token(TokenType.OPERATOR, c))
        else: # variable
            tokens.append(Token(TokenType.VARIABLE, c))
    return tokens

def priority_of_operator(operator):
    if operator == '+' or operator == '-':
        return 1
    elif operator == '*' or operator == '/':
        return 2
    else:
        raise ValueError(f"Unknown operator: {operator}")

def parse_tokens_recursive(tokens, left, right):
    count = right - left + 1
    if count <= 0:
        return None
    # check if it's single variable node
    if count == 1: # should be variable
        token = tokens[left] # left == right
        if token.type != TokenType.VARIABLE:
            raise TypeError("Expect VARIABLE token")
        return SyntaxNode(NodeType.VARIABLE, token.value)
    # check if it's a parentheses wrapped node
    if tokens[left].type == TokenType.LEFT_PARENTHESIS and tokens[right].type == TokenType.RIGHT_PARENTHESIS:
        root = parse_tokens_recursive(tokens, left + 1, right - 1)
        root.parentheses_wrapped = True
        return root

    # rightmost lowest priority op index
    op_index = None
    i = right
    while i >= left:
        token = tokens[i]
        if token.type == TokenType.RIGHT_PARENTHESIS:
            i = token.value - 1 # skip the whole parentheses
            continue
        elif token.type == TokenType.OPERATOR:
            if op_index:
                if priority_of_operator(token.value) < priority_of_operator(tokens[op_index].value):
                    op_index = i
            else:
                op_index = i
        i -= 1

    if not op_index:
        raise ValueError("No operator is found")
    left_subtree = parse_tokens_recursive(tokens, left, op_index - 1)
    right_subtree = parse_tokens_recursive(tokens, op_index + 1, right)
    root = SyntaxNode(NodeType.OPERATOR, tokens[op_index].value, left_subtree, right_subtree)
    return root

# from right to left, find the first of the operator of the lowest priority, set it as root
def parse_tokens_to_tree(tokens):
    return parse_tokens_recursive(tokens, 0, len(tokens) - 1)

def form_expression(root):
    if not root:
        return ""
    result = ""
    if root.parentheses_wrapped:
        result += "("
    result += form_expression(root.left) + str(root.value) + form_expression(root.right)
    if root.parentheses_wrapped:
        result += ")"
    return result

def parse_binding_to_context(binding):
    context = {}
    for item in binding.split(';'):
        pair = item.split('=')
        context[pair[0]] = int(pair[1])
    return context

def batch_bind(root, context):
    if not root:
        return
    if root.type == NodeType.VARIABLE:
        if not root.value in context:
            raise ValueError(f"Can't find value for variable {root.value}")
        num = context[root.value]
        root.type = NodeType.NUMBER
        root.value = num
    batch_bind(root.left, context)
    batch_bind(root.right, context)

def print_tree_recursive(root, level):
    if not root:
        return
    if root.parentheses_wrapped:
        print('  ' * level + '(')
    print('  ' * level + root.value)
    print_tree_recursive(root.left, level + 1)
    print_tree_recursive(root.right, level + 1)
    if root.parentheses_wrapped:
        print('  ' * level + ')')

def print_tree(root):
    print_tree_recursive(root, 0)

# if `should_batch_bind` is set to True, variables will be replaced with number value in the same round
def evaluate_expression(exp, binding, should_batch_bind):
    print(exp)
    tokens = parse_expression_to_tokens(exp)
    # print(tokens)
    root = parse_tokens_to_tree(tokens)
    # print_tree(root)
    # print(form_expression(root))
    context = parse_binding_to_context(binding)
    # print(context)
    if should_batch_bind:
        batch_bind(root, context)
        print(form_expression(root))
    root.evaluate(context, lambda: print(form_expression(root)))

def main():
    # exp = 'a+b-c'
    # exp = '(a+b)/c'
    exp = 'a/(b*(c-d))'
    binding = 'a=2;b=56;c=5;d=37'
    evaluate_expression(exp, binding, True)

main()