How to calculate an equation in a string, python

Hannibal The Vegan picture Hannibal The Vegan · Oct 9, 2015 · Viewed 9.5k times · Source

I have a variable that is function = '(2*1)+3'. How would I get it out of string form and calculate the answer? I tried using float(), int(float()) but I'm not sure if that's for numbers only or not.

Answer

mgilson picture mgilson · Oct 9, 2015

I've written this a couple times, and every time it seems that I lose the code...

A very simple (and "safe") calculator can be created using ast:

import ast
import operator

_OP_MAP = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.div,
    ast.Invert: operator.neg,
}


class Calc(ast.NodeVisitor):

    def visit_BinOp(self, node):
        left = self.visit(node.left)
        right = self.visit(node.right)
        return _OP_MAP[type(node.op)](left, right)

    def visit_Num(self, node):
        return node.n

    def visit_Expr(self, node):
        return self.visit(node.value)

    @classmethod
    def evaluate(cls, expression):
        tree = ast.parse(expression)
        calc = cls()
        return calc.visit(tree.body[0])


print Calc.evaluate('1 + 3 * (2 + 7)')

This calculator supports numbers, addition, subtraction, division, multiplication and negation (e.g. -6) and parenthesised groups. Order of operations are the same as Python which should be relatively intuitive... It can (almost trivially) be extended to support just about any unary or binary operator that python supports by adding the ast node type and corresponding operator/function to the _OP_MAP above.