Leetcode 736 Solution

#!/usr/bin/env python3


class Solution:
    def __init__(self):
        self.contexts = []

    def cut(self, expr):
        ans = []
        i = 1
        j = 1
        while j < len(expr) - 1:
            cnt = 0
            while j < len(expr) - 1 and (expr[j] != ' ' or cnt > 0):
                if expr[j] == '(':
                    cnt += 1
                elif expr[j] == ')':
                    cnt -= 1
                j += 1

            ans.append(expr[i:j])
            while j < len(expr) - 1 and expr[j] == ' ':
                j += 1
            i = j
        return ans

    def get_from_context(self, key):
        for context in reversed(self.contexts):
            if key in context:
                return context[key]
        return None

    def eval(self, expr):
        if expr[0] != '(':
            if expr.isdigit() or expr[0] == '-' and expr[1:].isdigit():
                return int(expr)
            else:
                return self.get_from_context(expr)

        tokens = self.cut(expr)
        op = tokens[0]

        if op == 'add':
            ans = 0
            for token in tokens[1:]:
                ans += self.eval(token)
            return ans
        elif op == 'mult':
            ans = 1
            for token in tokens[1:]:
                ans *= self.eval(token)
            return ans
        elif op == 'let':
            i = 1
            context = {}
            self.contexts.append(context)
            while tokens[i][0] != '(' and i + 1 < len(tokens):
                context[tokens[i]] = self.eval(tokens[i + 1])
                i += 2

            ans = None
            for j in range(i, len(tokens)):
                ans = self.eval(tokens[j])

            self.contexts.pop(-1)
            return ans

    def evaluate(self, expression: str) -> int:
        return self.eval(expression)