equation_evaluator.py

# -*- coding: cp1252 -*-
"""
Mathematical Expression Evaluator

Author: Sunjay Varma
Website: www.sunjay.ca

Supports a wide array of syntaxes and expressions.
* For complex numbers, use complex()
* use mode(RADIANS) or mode(DEGREES) to change trig mode
* You can define functions using 'f(x) =' notation (multivariable support e.g. 'abc(x, y, z) =')
* You can define variables with the 2 -> x notation (or more generally, expression -> expression)
* You can multiply using 2x and (2)(-4) notation
* Vectors are created for you and used as lists of numbers (YOU MUST HAVE THE 'Vector' MODULE.)
* 'fib' is an optional module that generates fibonacci numbers (use the fib function)
* The evaluator will balance parenthesis for you
* You can use |variable_name| for abs(variable_name)
* Use ||v|| to find the magnitude of Vector 'v'
* You can use ^ for exponents: e.g. f(x) = x^2 is the same as f(x) = x**2
* Technically, this supports ± (plus or minus) and square root symbols as well

You MUST set local to 'global_scope' when using meval if you want to have all the math functions!

It should stop you from trying any real Python expressions...let me know if you find any bugs!

General usage:
>>> meval('sin(pi/pi-1)', global_scope)
0
>>> meval('2x', {'x': 8})
16
>>> interact()
Type 'exit' or 'quit' to end this interpreter
Math> 2 + 2
4
Math> 2(4)
8
Math> 2(4)^2
32
Math> exit
"""

from __future__ import division
import cmath, math, re, code
from itertools import imap, izip
import keyword
import signal

# regular expressions
#r_term_parts = re.compile(r'(\d*)([^\^]+)?')
#r_plus_minus = re.compile(r'\+\s*\-')
# Used to get rid of all white space
r_white_space = re.compile(r'\s+')
# Math functions, not python functions: e.g. f(x) = y, abc(x, y, z) = q
r_func_def = re.compile(r'([a-zA-Z_]+)\(([^\)]+)\)\=')
# Function used for r_func_def to turn math functions into Python ones
repl_func_def = lambda m: m.group(1)+' = lambda ' + m.group(2) + ': '
# In math, we can do lots of fancy things with multiplication that we can't do in Python
# This translates most math multiplication notation into Python
r_terms_mul = re.compile(r'''                                         # fix things like 2x and 2(4)
                            (\d+                                      # number
                            |                                         # or
                            [^a-zA-Z\*\/+\-\(\,\=\:><\|!.; \t\n\r\f\v\xb1_]       # anything not a letter or operator
                            )
                            ([a-zA-Z_\(]+)          # either parenthesis or a letter
                            ''', re.VERBOSE)
# Catches functions with '__' around them
r_py_func = re.compile(r'__([^_]+)__')
# Supports assignment through 2 -> x
r_assign = re.compile(r'(.+)(?=\-\>)\-\>(.+)')
# Finds vectors
r_vector = re.compile(r'(?<!\w)(\([^\)]*\))')
# Finds angle functions
r_angle_func = re.compile(r'(a?sin|cos|tan{1}h?)(\([^\)]+\))')
# Supports mathematical abs notation: e.g. |x|
r_abs = re.compile(r'\|([^*/]+)\|')
# Finds the magnitude of a vector ||v||
r_mag = re.compile(r'\|\|([^*/]+)\|\|')
#r_mag_mul = re.compile(r'(\|{1,2})(.+\|)\1(\|{1,2})(.+\|)\3')
# Finds integers
r_int = re.compile(r'(?<!\.)(\d+)(?!\.)')
# Finds invalid integers (syntax error)
r_invalid_int = re.compile(r'(\d+\.\d+)\.0')
# Finds decimal numbers
r_deg = re.compile(r"[+-]?((\d+(\.\d*)?)|\.\d+)([eE][+-]?[0-9]+)?\xb0")
# Looks for words (variables, function names, etc)
r_word = re.compile(r"\w+")
# Looks for the square root symbol
r_sqrt = re.compile(u"\u221a"+r"(\()?")

# constants
DEFAULT_VARIABLE = "x"
# calculator modes
RADIAN = 'RADIAN'
DEGREE = 'DEGREE'
MODE_CONST = 'CALC_MODE' # used for detecting modes
DEFAULT_MODE = RADIAN

# errors
class Error(Exception): pass # base exception for this module
class PotentialSecurityThreat(Error): pass
class CalculationFailure(Error): pass

# create a global scope which will be used during the evaluation of the expression
global_scope = {
    '__builtins__': {},
    MODE_CONST: DEFAULT_MODE,
    'RADIAN': RADIAN,
    'DEGREE': DEGREE
}

# make sure the user has access to math and cmath variables
for mod in [math, cmath]:
    global_scope.update(vars(mod))

# things that are needed from the builtin functions
builtins_needed = {
    'abs': abs,
    'int': int,
    'float': float,
    'round': round,
}

# add specific built in functions (and nothing else)
global_scope.update(builtins_needed)

# load extension modules
for name in 'Vector fib'.split():
    try:
        global_scope[name] = value = getattr(__import__(name.lower()), name)
        exec "global %s"%name
        exec "%s = getattr(__import__(%r.lower()), %r)"%(name, name, name)
    except:
        pass

def vrange(start, stop=None, step=1):
    """Vector of range(start, stop, step)"""
    if stop is None:
        stop = start
        start = 0
    return Vector(range(int(start), int(stop), int(step)))

# add this to the global scope so the user has access
global_scope["vrange"] = vrange

def dextend(*args):
    """Update the first argument with all the others"""
    ext = args[0]
    for d in args[1:]:
        ext.update(d)
    return ext

def make_mode_func(scope):
    """Make a function to define the calculator mode"""
    def _mode(mode=None):
        if mode and mode in [RADIAN, DEGREE]:
            scope[MODE_CONST] = mode
        return scope[MODE_CONST]
    scope['mode'] = _mode

def sstartswith(string, prefixes):
    """Return True if string starts with any of the prefixes"""
    for prefix in prefixes:
        if string.startswith(prefix):
            return True

def sreplace(s, old, new):
    """Basically goes through old and replaces old[i] with the corresponding new[i]"""
    for x, y in izip(old, new):
        s = s.replace(x, y)
    return s

BRACKETS = dict(izip('[{(', ']})'))

def balance(s):
    """
    Balance the parentheses within a string

    sin(cos(tan(x) --> sin(cos(tan(x)))
    """
    openedb = BRACKETS.keys()
    closedb = BRACKETS.values()
    if [s.count(x) for x in openedb] == [s.count(x) for x in closedb]:
        return s # balanced
    open_br = []
    for c in s:
        if c in openedb:
            open_br.append(c)
        elif c in closedb:
            if not open_br or not BRACKETS[open_br[-1]] == c:
                    raise SyntaxError('Invalid Syntax!') # too many closing brackets
            open_br.pop()
    for br in open_br:
        s += BRACKETS[br]
    return s

def check_security_threat(expr, orig):
    """Looks for potential security threats by seeking out python keywords and underscores"""
    if sstartswith(expr, keyword.kwlist):
        pass # any python key word
    # double under score functions
    elif next(r_py_func.finditer(expr), None):
        pass
    else:
        return
    # If any of the tests pass, we probably have a security threat
    raise PotentialSecurityThreat(orig)

def mfix(expr):
    """fixes up an expression or equation for evaluation"""
    # Save the original expression for accurate error reporting
    orig_expr = expr

    # Take out all white space
    expr = r_white_space.sub("", expr)
    if not expr.strip():
        return ''

    # checks for different kinds of potential security threats
    check_security_threat(expr, orig_expr)

    #expr = r_sqrt.sub("sqrt(", expr).encode('UTF-8')
    # Balance all the different types of parenthesis
    expr = balance(expr)
    # vector and absoulute operators
    #expr = r_mag_mul.sub(lambda m: m.group(1)+m.group(2)+m.group(1)+'*'+m.group(3)+m.group(4)+m.group(3), expr)

    # Get vector magnitudes
    expr = r_mag.sub(lambda m: m.group(1)+'.magnitude', expr)
    # Change |x| to abs(x)
    expr = r_abs.sub(lambda m: 'abs('+m.group(1)+')', expr)
    # Replace ^ with **, )*(, and remove quotation marks
    # This is just to save space and avoid repeating the same code
    expr = sreplace(expr, ['^', ')(', '"', "'"], ['**', ')*(', '', ''])
    # Replace all different types of brackets with parenthesis
    expr = sreplace(expr, '{}[]', '()()')
    # Replace 2 -> x with x = 2
    expr = r_assign.sub(lambda m: m.group(2)+'='+m.group(1), expr)
    # Change a(x) = y into a Python function
    expr = r_func_def.sub(repl_func_def, expr)
    # Change things like 2x to 2*x
    expr = r_terms_mul.sub(lambda m: m.group(1)+"*"+m.group(2), expr)

    # Account for the plus or minus symbol
    if '\xb1' in expr: # plus or minus symbol
        # Check if this is an assignment statement
        if '=' in expr.replace('==', ''):
            # Get rid of other equal signs
            expr = expr.replace('==', '__HOLD__')
            if 'lambda' in expr:
                i = expr.index(':', expr.index('lambda'))
                # Find the assignment statement
                i_eq = expr.index('=')
                name, extra, value = expr.split('=')[0], expr[i_eq+1:i+1], expr[i+1:]
            else:
                name, value, extra = expr.split('=') + ['']
            name += '='
            name, value = map(lambda x: x.replace('__HOLD__', '=='), [name, value])
        else:
            name, extra, value = '', '', expr

        # Get both the positive and negative value
        expr = name + extra + 'Vector(' + value.replace('\xb1', '+') + ',' + value.replace('\xb1', '-') + ')'

    # Look for vectors
    for mobj in r_vector.finditer(expr):
        vec = mobj.group(1)
        if ',' in vec or vec == '()':
            expr = expr[:mobj.start()] + 'Vector' + vec + expr[mobj.end():]

    # force float
    expr = r_int.sub(lambda m: m.group(1)+'.0', expr)
    # fix invalid integers created by this
    expr = r_invalid_int.sub(lambda m: m.group(1), expr)
    #expr = re.sub(r'([[{(])\*([[{(])', lambda m: m.group(1)+m.group(2), expr)
    return expr

def mcompile(expr):
    """Fixes up a math expression and compiles into Python code"""
    expr = mfix(expr)
    if "=" in expr.replace('==', ''):
        sym = "exec"
    else:
        sym = "eval"
    return code.compile_command(expr, symbol=sym)

def _calc_fail(signum, frame):
    raise CalculationFailure('Calculation Failed!')

# Used to reformat results
class SimpleRepr:
    """Give simple repr'd values for objects"""
    def __init__(self, data):
        self.data = data

    def __repr__(self):
        return self.data

def find_name(x, scope):
    for name, value in scope.iteritems():
        if x == value:
            return name

LAMBDA_NAME = (lambda: None).__name__

def reformat_result(iterable, scope={}):
    """Change the representation of items in result such as callable or floating point objects"""
    new = []
    for x in iterable:
        if hasattr(x, '__call__'):
            name = getattr(x, '__name__')
            if name == LAMBDA_NAME:
                name = find_name(x, scope)
            new.append(SimpleRepr(name))
        elif hasattr(x, '__iter__'):
            new.append(reformat_names(x))
        elif isinstance(x, complex):
            if not x.imag:
                x = float(x.real)
            new.append(x)
        else:
            # float which could be an int in simplest form
            if hasattr(x, 'is_integer') and x.is_integer():
                x = int(x)
            new.append(x)
    return tuple(new)

def meval(expr, local={}):
    """
    Evaluate an expression

    Note: You will NEED to set local to global_scope in order to use all the math functions
    """
    # the expression is either a function or a variable in the local scope
    if mfix(expr) in local:
        obj = local[mfix(expr)]
        # if obj is a function
        if hasattr(obj, '__call__'):
            return expr
        # otherwise, return the object itself
        else:
            return obj

    if MODE_CONST in local and local[MODE_CONST] == DEGREE:
        # m.group(2) has parenthesis already around it
        expr = r_angle_func.sub(lambda m: m.group(1)+'(radians'+m.group(2)+')', expr)

    # The signal code here makes sure that the execution doesn't take too long
    if hasattr(signal, "SIGALRM"):
        signal.signal(signal.SIGALRM, _calc_fail)
        signal.alarm(5)
    # Evaluate the result
    result = eval(mcompile(expr), local, local)
    if hasattr(signal, "SIGALRM"):
        signal.alarm(0)
    if isinstance(result, complex):
        if not result.imag:
            result = float(result.real)
        else:
            # Display complex numbers properly
            result = "%s+%si"%(result.real, result.imag)
    # float which could be an int in simplest form
    if hasattr(result, 'is_integer') and result.is_integer():
        result = int(result)
    elif hasattr(result, '__iter__'):
        if isinstance(result, set):
            result = tuple(result)
        result = reformat_result(result, local)
    return result

def get_line():
    """Gets a line of input"""
    try:
        return raw_input('Math> ')
    except KeyboardInterrupt:
        print 'KeyboardInterrupt'
    except:
        # Obviously must be in scope to be used
        print_exc()
    return ""

def interact(scope=global_scope.copy()):
    """Creates an interactive Math prompt and gives access to default math functions"""
    from traceback import print_exc

    print "Type 'exit' or 'quit' to end this interpreter"

    # Allow for control of the trignometric function modes
    if not 'mode' in scope:
        make_mode_func(scope)

    # Get a single line of input
    line = get_line()
    while True:
        # Let the user quit
        if line in ('exit', 'quit'):
            break
        try:
            for expr in line.split(';'):
                if not expr:
                    continue
                result = meval(expr, scope)
                if result is not None:
                    print result
                    # Set the previously found result to the variable _
                    scope['_'] = result
        except KeyboardInterrupt:
            print 'KeyboardInterrupt'
        except:
            print_exc()
        line = get_line()

if __name__ == "__main__":
    interact()