calc.py

import re
import cmath
import math
import random
import time
import operator
from copy import deepcopy

epsilon = 1e-10

# http://stackoverflow.com/a/2182437
class Enum (set):
	def __init__ (self, *args):
		super (Enum, self).__init__ (args)

	def __getattr__ (self, name):
		if name in self:
			return name
		raise AttributeError

	def __setattr__ (self, name, value):
		raise AttributeError

class Operator (object):
	def __init__ (self, name, symbol, operands, priority, function):
		self.name = name
		self.symbol = symbol
		self.operands = operands
		self.priority = priority
		self.function = function

	def __str__ (self):
		return '''<operator %s>''' % self.name

	def __repr__ (self):
		return self.__str__()

class FunctionCall (object):
	def __init__ (self, funcname, args):
		assert (type(args) is list)
		self.funcname = funcname
		self.args = args

	def __str__ (self):
		return '''<function %s (%s)>''' % (self.funcname, self.args)

	def __repr__ (self):
		return self.__str__()

def do_realf (func, *args):
	for x in args:
		if x.imag:
			raise ValueError ('%s called with a complex number' % func.__name__)

	return func (*[x.real for x in args])

def do_intf (func, *args):
	for x in args:
		if x.imag:
			raise ValueError ('%s called with a complex number' % func.__name__)

		if x.real - math.floor (x.real):
			raise ValueError ('%s called with a floating point number' % func.__name__)

	return func (*[int (x.real) for x in args])

def realf (func):
	return lambda *args: do_realf (func, *args)

def intf (func):
	return lambda *args: do_intf (func, *args)

Operators = []
OperatorData = {
	'lneg':		{ 'symbol': '!',  'operands': 1, 'priority': 5, 'function': lambda x: not x },
	'compl':	{ 'symbol': '~',  'operands': 1, 'priority': 5, 'function': intf (operator.inv) },
	'neg':		{ 'symbol': '-',  'operands': 1, 'priority': 5, 'function': lambda x: -x },
	'pow':		{ 'symbol': '**', 'operands': 2, 'priority': 10, 'function': lambda x, y: x ** y },
	'mul':		{ 'symbol': '*',  'operands': 2, 'priority': 50, 'function': lambda x, y: x * y },
	'div':		{ 'symbol': '/',  'operands': 2, 'priority': 50, 'function': lambda x, y: x / y },
	'mod':		{ 'symbol': '%',  'operands': 2, 'priority': 50, 'function': lambda x, y: math.fmod (x, y) },
	'add':		{ 'symbol': '+',  'operands': 2, 'priority': 100, 'function': lambda x, y: x + y },
	'sub':		{ 'symbol': '-',  'operands': 2, 'priority': 100, 'function': lambda x, y: x - y },
	'eq':		{ 'symbol': '==', 'operands': 2, 'priority': 500, 'function': lambda x, y: x == y },
	'neq':		{ 'symbol': '!=', 'operands': 2, 'priority': 500, 'function': lambda x, y: x != y },
	'lt':		{ 'symbol': '<',  'operands': 2, 'priority': 500, 'function': lambda x, y: x < y },
	'lteq':		{ 'symbol': '<=', 'operands': 2, 'priority': 500, 'function': lambda x, y: x <= y },
	'gt':		{ 'symbol': '>',  'operands': 2, 'priority': 500, 'function': lambda x, y: x > y },
	'gteq':		{ 'symbol': '>=', 'operands': 2, 'priority': 500, 'function': lambda x, y: x >= y },
	'btand':	{ 'symbol': '&',  'operands': 2, 'priority': 600, 'function': intf (operator.and_) },
	'btxor':	{ 'symbol': '^',  'operands': 2, 'priority': 601, 'function': intf (operator.xor) },
	'btor':		{ 'symbol': '|',  'operands': 2, 'priority': 602, 'function': intf (operator.or_) },
	'and':		{ 'symbol': '&&', 'operands': 2, 'priority': 603, 'function': lambda x, y: x and y },
	'or':		{ 'symbol': '||', 'operands': 2, 'priority': 604, 'function': lambda x, y: x or y },
}

for name, data in OperatorData.iteritems():
	Operators.append (Operator (name=name, symbol=data['symbol'], operands=data['operands'],
		priority=data['priority'], function=data['function']))

OperatorSymbols={}
for op in Operators:
	if op.symbol in OperatorSymbols:
		OperatorSymbols[op.symbol].append (op)
	else:
		OperatorSymbols[op.symbol] = [op]

Constants = {
	'pi': cmath.pi,
	'e': cmath.e,
	'phi': 1.6180339887498948482,
	'epsilon': epsilon,
}

random.seed (time.time())

Functions = {
	'round':	{ 'function': realf (round) },
	'floor':	{ 'function': realf (math.floor) },
	'ceil':		{ 'function': realf (math.ceil) },
	'fact':		{ 'function': intf (math.factorial) },
	'abs':		{ 'function': realf (math.fabs) },
	'degrees':	{ 'function': realf (math.degrees) },
	'radians':	{ 'function': realf (math.radians) },
	'erf':		{ 'function': realf (math.erf) },
	'erfc':		{ 'function': realf (math.erfc) },
	'gamma':	{ 'function': realf (math.gamma) },
	'lgamma':	{ 'function': realf (math.lgamma) },
	'sqrt':		{ 'function': cmath.sqrt },
	'ln':		{ 'function': cmath.log },
	'log':		{ 'function': cmath.log10 },
	'sin':		{ 'function': cmath.sin },
	'sinh':		{ 'function': cmath.sinh },
	'asin':		{ 'function': cmath.asin },
	'asinh':	{ 'function': cmath.asinh },
	'cos':		{ 'function': cmath.cos },
	'cosh':		{ 'function': cmath.cosh },
	'acos':		{ 'function': cmath.acos },
	'acosh':	{ 'function': cmath.acosh },
	'tan':		{ 'function': cmath.tan },
	'tanh':		{ 'function': cmath.tanh },
	'atan':		{ 'function': cmath.atan },
	'atanh':	{ 'function': cmath.atanh },
	'exp':		{ 'function': cmath.exp },
	'phase':	{ 'function': cmath.phase },
	'lg':		{ 'function': lambda x: cmath.log (x, 2) },
	're':		{ 'function': lambda x: x.real },
	'im':		{ 'function': lambda x: x.imag },
	'conjugate':{ 'function': lambda x: x.conjugate() },
	'rand':		{ 'function': random.random },
}

Tokens = ['(', ')']

# Symbol table
SymbolType = Enum ('CONSTANT', 'FUNCTION', 'OPERATOR', 'TOKEN')
SymbolTypes = {}
Symbols = []

for name, value in Constants.iteritems():
	Symbols.append (name)
	SymbolTypes[name] = SymbolType.CONSTANT

for name, data in Functions.iteritems():
	Symbols.append (name)
	SymbolTypes[name] = SymbolType.FUNCTION

for op in Operators:
	if op.symbol not in Symbols:
		Symbols.append (op.symbol)
	SymbolTypes[op.symbol] = SymbolType.OPERATOR

for name in Tokens:
	SymbolTypes[name] = SymbolType.TOKEN

Symbols += Tokens
Symbols = sorted (Symbols, key=lambda x: len (x), reverse=True)

def parse_number (expr):
	"""Tries to parse a number from the expression. Returns (value, length) on success."""
	i = 0
	value = None

	if expr[0] == 'i':
		# Special case, we just have 'i' -- need special handling here because otherwise this would
		# call float('') in the complex number routine, which throws an error.
		# TODO: maybe 'i' can be a symbol instead?
		return (1j, 1)

	# Try parse increasingly long substrings until we are unable to create a float or complex number
	# from it.
	try:
		while i < len (expr):
			if expr[i] == 'i':
				value = complex (0.0, float (expr[:i]))
			else:
				value = complex (float (expr [:i + 1]), 0.0)
			i += 1

		return (value, i)
	except ValueError:
		if i != 0:
			# Got a number (the error happened when we tried to parse past the number)
			return (value, i)
		else:
			# The error happened on the first character. So this is not a number.
			return None

def parse_symbol (expr):
	for sym in Symbols:
		if expr[:len (sym)] == sym:
			return sym

	return None

def tokenize (expr):
	expr = re.sub ('\s+', '', expr.strip())
	i=0
	tokens = []

	while i < len(expr):
		sym = parse_symbol (expr[i:])

		if sym:
			symtype = SymbolTypes[sym]
			if symtype == SymbolType.CONSTANT:
				tokens.append (Constants[sym])
			else:
				tokens.append (sym)

			i += len(sym)
			continue

		result = parse_number (expr[i:])
		if result:
			num, length = result
			tokens.append (num)
			i += length
			continue

		raise ValueError ("""bad expression, couldn't parse: %s""" % expr[i:])

	return tokens

def rindex (li, value):
	return (len(li) - 1) - li[::-1].index(value)

def process_parens (expr):
	"""Processes parentheses of expr into sublists in-place.
	[1.0, '*', '(', 3.5, '+', 1j, ')']
	-> [1.0, '*', [3.5, '+', 1j]]"""
	if '(' not in expr and ')' not in expr:
		return

	try:
		parenStart = rindex (expr, '(')
		parenEnd = expr.index (')', parenStart)
	except ValueError:
		raise ValueError ("""mismatched parentheses in expression: %s""" % expr)

	subexpr = expr[parenStart + 1:parenEnd]
	del expr[parenStart + 1:parenEnd + 1]
	expr[parenStart] = subexpr
	process_parens (expr)

def process_functions (expr):
	"""Processes functions in-place"""
	i = 0
	while i < len (expr):
		if type (expr[i]) is list:
			process_functions (expr[i])

		if (type(expr[i]) is not str) or (expr[i] not in Functions):
			i += 1
			continue

		# Ensure that arguments follow
		if (i + 1 >= len (expr)) or (type (expr[i + 1]) is not list):
			raise ValueError ("""function %s used without arguments""" % expr[i])

		args = expr[i + 1]
		del expr[i + 1]
		expr[i] = FunctionCall (expr[i], args)
		i += 1

def is_operand (x):
	# Operands can be either lists (which also mean parens, thus can be single-expressions)
	# or complex numbers
	return type(x) in [complex, list]

def find_fitting_operator (sym, numOperands):
	# Pass 1: exact numOperands match
	for op in Operators:
		if op.symbol != sym:
			continue

		if op.operands == numOperands:
			# print '%s, %d -> %s (1st pass)' % (sym, numOperands, op)
			return op

	# Pass 2: by symbol
	for op in Operators:
		if op.symbol == sym:
			# print '%s, %d -> %s (2nd pass)' % (sym, numOperands, op)
			return op

	raise ValueError ('''unknown operator %s!''' % sym)

def process_operators (expr):
	"""Processes operators"""
	i = 0

	# Find all operators in this expression
	while i < len (expr):
		if type (expr[i]) is list:
			process_functions (expr[i])
			process_operators (expr[i])

		if type (expr[i]) is FunctionCall:
			process_functions (expr[i].args)
			process_operators (expr[i].args)

		if (type(expr[i]) is not str) or (expr[i] not in OperatorSymbols):
			i += 1
			continue

		args = []
		argIndices = []
		if i > 0 and is_operand (expr[i - 1]):
			args.append (expr[i - 1])
			argIndices.append (i - 1)

		if i - 1 < len(expr) and is_operand (expr[i + 1]):
			args.append (expr[i + 1])
			argIndices.append (i + 1)

		# Resolve operators with the same symbol based on operand count
		numOperands = 0
		for arg in args:
			if is_operand (arg):
				numOperands += 1

		expr[i] = find_fitting_operator (expr[i], numOperands)
		i += 1

def find_priority_operator (expr):
	"""Finds the operator with most priority in the expression"""
	bestOp = None
	bestOpIndex = -1

	for i in range (0, len (expr)):
		sym = expr[i]

		if type (sym) is not Operator:
			continue

		if not bestOp or sym.priority < bestOp.priority:
			bestOp = sym
			bestOpIndex = i

	return (bestOp, bestOpIndex)

def realPrint (x):
	print x

tabs=''
def evaluate (expr, verbose=False):
	global tabs
	printFunc = realPrint if verbose else lambda x:None
	printFunc (tabs + 'Preprocess: %s' % expr)

	# If there are sub-expressions in here, those need to be evaluated first
	i = 0
	while i < len (expr):
		sym = expr[i]

		if type (sym) is list and sym:
			printFunc (tabs + 'Evaluating sub-expression: %s' % (sym))
			tabs += '\t'
			expr[i] = evaluate (list (sym), verbose)
			tabs = tabs[:-1]
			printFunc (tabs + '-> %s' % expr[i])

		# If there are function calls, evaluate those
		if type (sym) is FunctionCall:
			tabs += '\t'
			if sym.args:
				sym.args = [evaluate (sym.args, verbose)]
			tabs = tabs[:-1]

			printFunc (tabs + 'Evaluating function call: %s' % (sym))
			expr[i] = Functions[sym.funcname]['function'] (*sym.args)
			printFunc (tabs + '-> %s' % expr[i])
			
		i += 1

	printFunc (tabs + 'Evaluate: %s' % expr)
	runaway = 0
	while True:
		runaway += 1
		if runaway > 1000:
			raise ValueError ('infinite loop detected')

		op, i = find_priority_operator (expr)
		if not op:
			break

		if op.operands == 2:
			argIndices = [i - 1, i + 1]
		else:
			argIndices = [i + 1]

		args = [expr[x] for x in argIndices]
		argIndices = sorted (argIndices, reverse=True)
		printFunc (tabs + 'Processing: (%s, %d) with args %s' % (op, i, args))
		expr[i] = op.function (*args)
		printFunc (tabs + '-> %s' % expr[i])

		for i2 in argIndices:
			del expr[i2]

	printFunc (tabs + 'Result: %s' % expr[0])

	if len(expr) != 1:
		printFunc (tabs + 'Bad expression detected, tokens: %s' % expr)
		raise ValueError ('malformed expression')

	return expr[0]

def repr_number (x):
	"""Returns a string representation for a real number"""
	rep = '%.10f' % x

	if '.' in rep:
		while rep[-1] == '0':
			rep = rep[:-1]

		if rep[-1] == '.':
			rep = rep[:-1]

	return rep

def repr_imaginary (x):
	rep = repr_number (x)

	if rep == '1':
		return 'i'

	if rep == '-1':
		return '-i'

	return rep + 'i'

def represent (x):
	"""Returns a string representation of a float or complex number"""
	if math.fabs (x.imag) > epsilon:
		if math.fabs (x.real) > epsilon:
			# Complex number
			return '%s %s %s' % (repr_number (x.real),
				'+' if x.imag >= 0 else '-',
				repr_imaginary (math.fabs (x.imag)))
		else:
			# Pure imaginary number
			return repr_imaginary (x.imag)
	else:
		# Real number
		return repr_number (x.real)

def calc (expr, verbose=False):
	expr = tokenize (expr)
	process_parens (expr)
	process_functions (expr)
	process_operators (expr)
	return represent (evaluate (expr, verbose))