diff --git a/frappy/lib/mathparser.py b/frappy/lib/mathparser.py new file mode 100644 index 00000000..52cc03bd --- /dev/null +++ b/frappy/lib/mathparser.py @@ -0,0 +1,90 @@ +# ***************************************************************************** +# +# This program is free software; you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the Free Software +# Foundation; either version 2 of the License, or (at your option) any later +# version. +# +# This program is distributed in the hope that it will be useful, but WITHOUT +# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS +# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more +# details. +# +# You should have received a copy of the GNU General Public License along with +# this program; if not, write to the Free Software Foundation, Inc., +# 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA +# +# Module authors: +# Markus Zolliker +# Anik Stark +# https://stackoverflow.com/questions/43836866/safely-evaluate-simple-string-equation +# +# ***************************************************************************** + +import math +import ast +import operator as op + +class MathParser: + _operators2method = { + ast.Add: op.add, + ast.Sub: op.sub, + ast.BitXor: op.xor, + ast.Or: op.or_, + ast.And: op.and_, + ast.Mod: op.mod, + ast.Mult: op.mul, + ast.Div: op.truediv, + ast.Pow: op.pow, + ast.FloorDiv: op.floordiv, + ast.USub: op.neg, + ast.UAdd: lambda a:a} + + def __init__(self, math=True, **kwargs): + self._vars = kwargs + if not math: + self._alt_name = self._no_alt_name + + def _Name(self, name): + try: + return self._vars[name] # look up in user-provided dict + except KeyError: + return self._alt_name(name) # fall back to math functions + + @staticmethod + def _alt_name(name): + if name.startswith("_"): # prevent access to hidden names + raise NameError(f"{name!r}") + try: + return getattr(math, name) + except AttributeError: + raise NameError(f"{name!r}") + + @staticmethod + def _no_alt_name(name): + raise NameError(f"{name!r}") + + def eval_(self, node): + if isinstance(node, ast.Expression): + return self.eval_(node.body) + if isinstance(node, ast.Constant): # return the number + return node.value + if isinstance(node, ast.Name): # return variable or math function + return self._Name(node.id) + if isinstance(node, ast.BinOp): # evaluate binary operations + method = self._operators2method[type(node.op)] + return method( self.eval_(node.left), self.eval_(node.right)) + if isinstance(node, ast.UnaryOp): # handle operators + method = self._operators2method[type(node.op)] + return method( self.eval_(node.operand) ) + if isinstance(node, ast.Attribute): # handle attributes (e.g. math.cos) + return getattr(self.eval_(node.value), node.attr) + if isinstance(node, ast.Call): # evaluate the function and its arguments, calls function + return self.eval_(node.func)( + *(self.eval_(a) for a in node.args), + **{k.arg:self.eval_(k.value) for k in node.keywords}) + raise TypeError(node) + + def calculate(self, expr, **kwargs): + self._vars.update(kwargs) + return self.eval_(ast.parse(expr, mode='eval'))