通过ast实现对一个数学表达式的字符串的计算

当我们有一个如“3+4*sin(pi)/log(2)”的一个字符串，需要将其进行计算。首先要考虑其优先级，在对中如“sin()”、“log()”等特殊计算进行解析，这时可以考虑ast。

将字符解析为树

import ast
import math

def calculate(expression):
    expression = expression.replace('^', str("**")) #python的次方为**
    root = ast.parse(expression, mode='eval') #通过ast，将字符串解析为树
    calculator = Calculator()
    # print(ast.dump(root))
    calculator.visit(root)
    return calculator.result

对树进行计算

class Calculator(ast.NodeVisitor):
    def __init__(self):
        self.result = None

    def visit_Num(self, node):
        self.result = node.n

    def visit_Name(self, node):
        if node.id == 'e':
            self.result = math.e
        elif node.id == 'pi':
            self.result = math.pi

    def visit_BinOp(self, node):
        self.visit(node.left)
        left_value = self.result
        self.visit(node.right)
        right_value = self.result
        operator = node.op
        if isinstance(operator, ast.Add):
            self.result = left_value + right_value
        elif isinstance(operator, ast.Sub):
            self.result = left_value - right_value
        elif isinstance(operator, ast.Mult):
            self.result = left_value * right_value
        elif isinstance(operator, ast.Div):
            self.result = left_value / right_value
        elif isinstance(operator, ast.FloorDiv):
            self.result = left_value // right_value
        elif isinstance(operator, ast.Mod):
            self.result = left_value % right_value
        elif isinstance(operator, ast.Pow):
            self.result = left_value ** right_value

    def visit_Call(self, node):
        func_name = node.func.id
        if func_name == 'sin':
            self.visit(node.args[0])
            arg_value = self.result
            self.result = math.sin(arg_value)
        elif func_name == 'cos':
            self.visit(node.args[0])
            arg_value = self.result
            self.result = math.cos(arg_value)
        elif func_name == 'tan':
            self.visit(node.args[0])
            arg_value = self.result
            self.result = math.tan(arg_value)
        elif func_name == 'sqrt':
            self.visit(node.args[0])
            arg_value = self.result
            self.result = math.sqrt(arg_value)
        elif func_name == 'log':
            # 检查参数个数
            arg_count = len(node.args)
            if arg_count == 1:
                # 计算自然对数
                self.visit(node.args[0])
                arg_value = self.result
                self.result = math.log(arg_value)
            elif arg_count == 2:
                # 计算以指定底数的对数
                self.visit(node.args[0])
                arg_value = self.result
                self.visit(node.args[1])
                base_value = self.result
                self.result = math.log(arg_value, base_value)