frappy/lib: add math parser to evaluate a string function

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 <markus.zolliker@psi.ch>
+#   Anik Stark <anik.stark@psi.ch>
+#   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'))