# HG changeset patch
# User Teemu Piippo <crimsondusk64@gmail.com>
# Date 1444077344 -10800
# Node ID 6e7cb38f248d68fc770ba5ca56c74b34a19e848d
# Parent  0ee4ff9137478769d4cc859dc52dc07e89f6197e
Calculator now supports variables

diff -r 0ee4ff913747 -r 6e7cb38f248d calculator.py
--- a/calculator.py	Mon Oct 05 21:55:53 2015 +0300
+++ b/calculator.py	Mon Oct 05 23:35:44 2015 +0300
@@ -35,18 +35,20 @@
 import operator
 import string
 import enum
+from fraction import Fraction
 from copy import deepcopy
 from math import pi as π
 
 ε = 1e-10
 
 class Operator (object):
-	def __init__ (self, name, symbol, operands, priority, function):
+	def __init__ (self, name, symbol, operands, priority, function, assign):
 		self.name = name
 		self.symbol = symbol
 		self.operands = operands
 		self.priority = priority
 		self.function = function
+		self.assign = assign
 
 	def __repr__ (self):
 		return '<operator %s>' % (self.name)
@@ -64,11 +66,27 @@
 
 class Value (object):
 	def __init__ (self, value):
+		if isinstance (value, Fraction) and value.denominator == 1:
+			value = value.numerator
 		self.value = value
 
 	def __repr__ (self):
 		return '''Value(%r)''' % (self.value)
 
+class Name (object):
+	def __init__ (self, name):
+		self.name = name
+
+	def __repr__ (self):
+		return '''Name(%r)''' % (self.name)
+
+class AssignmentResult (object):
+	def __init__ (self, name, value):
+		self.name, self.value = name, value
+
+	def __repr__ (self):
+		return '''AssignmentResult(%r, %r)''' % (self.name, self.value)
+
 def do_realf (func, *args):
 	for x in args:
 		if x.imag:
@@ -109,10 +127,20 @@
 def factorial (x):
 	return math.gamma (x + 1)
 
+def is_int (x):
+	return math.fabs (x - math.floor(x)) < ε
+
+def div(a,b):
+	if is_int(a) and is_int(b):
+		return Fraction(a,b)
+	else:
+		return a / b
+
+def assignment (name, value, calculator):
+	calculator.variables[name.name] = value.value
+
 Operators = {}
 OperatorData = {
-	'sci':		{ 'symbol': 'e',  'operands': 2, 'priority': 1, 'function': intf (scientific) },
-	'dice':		{ 'symbol': 'd',  'operands': 2, 'priority': 2, 'function': intf (dice) },
 	'not':		{ '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 },
@@ -121,7 +149,7 @@
 	'fact':		{ 'symbol': '!',  'operands': -1, 'priority': 8, 'function': realf(factorial) },
 	'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 },
+	'div':		{ 'symbol': '/',  'operands': 2, 'priority': 50, 'function': div },
 	'mod':		{ 'symbol': '%',  'operands': 2, 'priority': 50, 'function': realf (math.fmod) },
 	'add':		{ 'symbol': '+',  'operands': 2, 'priority': 100, 'function': lambda x, y: x + y },
 	'sub':		{ 'symbol': '-',  'operands': 2, 'priority': 100, 'function': lambda x, y: x - y },
@@ -136,6 +164,7 @@
 	'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 },
+	'assign':	{ 'symbol': '=',  'operands': 2, 'priority': 999, 'function': assignment, 'assign':True },
 }
 
 def powerfunction (n):
@@ -143,14 +172,15 @@
 
 superscripts='⁰¹²³⁴⁵⁶⁷⁸⁹'
 for i, sup in enumerate(superscripts):
-	print (i, sup)
 	func = powerfunction(i)
 	func.__name__ = 'x' + sup
 	OperatorData['sup' + str(i)] = { 'symbol': sup, 'operands': -1, 'priority': 3, 'function': func }
 
 for name, data in OperatorData.items():
-	Operators[name] = Operator (name=name, symbol=data['symbol'], operands=data['operands'],
-		priority=data['priority'], function=data['function'])
+	if 'assign' not in data:
+		data['assign'] = False
+
+	Operators[name] = Operator (name=name, **data)
 
 OperatorSymbols={}
 for op in Operators.values():
@@ -240,10 +270,6 @@
 	'ε': 'epsilon',
 	'∨': '||',
 	'∧': '&&',
-	'⊻': '^',
-	'ℜ': 're',
-	'ℑ': 'im',
-	'fact': 'factorial',
 }
 
 for name, value in Constants.items():
@@ -287,11 +313,17 @@
 def is_int (x):
 	return math.fabs (x - math.floor(x)) < ε
 
+identifierchars = set(string.ascii_lowercase + string.ascii_uppercase + string.digits + '_')
+def is_identifier(x):
+	return x and (x[0] not in string.digits) and (set(x) - identifierchars == set())
+
 class Calculator (object):
 	def __init__ (self, verbose=False):
 		self.previous_result = None
 		self.preferred_base = None
 		self.verbose = verbose
+		self.variables = {}
+		self.namepattern = re.compile (r'^([A-Za-z_][A-Za-z0-9_]*)')
 
 	def set_preferred_base (self, x):
 		self.preferred_base = x
@@ -428,7 +460,14 @@
 				i += length
 				continue
 
-			raise ValueError ("""bad expression, couldn't parse: %s""" % expr[i:])
+			match = self.namepattern.match (expr[i:])
+			if match:
+				name = match.group(1)
+				i += len(name)
+				tokens.append (Name(name))
+				continue
+
+			raise ValueError ("""couldn't parse: %s""" % expr[i:])
 
 		return tokens
 
@@ -470,26 +509,36 @@
 			expr[i] = FunctionCall (expr[i], args)
 			i += 1
 
+	def process_variables (self, expr):
+		for i, x in enumerate(expr):
+			if isinstance(x, list):
+				self.process_variables(x)
+			elif isinstance(x, Name):
+				try:
+					expr[i].value = Value(self.variables[x.name])
+				except KeyError as e:
+					pass
+
 	def is_operand (self, x):
 		# Operands can be either lists (which also mean parens, thus can be single-expressions)
 		# or complex numbers
-		return type(x) in [list, FunctionCall, Value]
+		return type(x) in [list, FunctionCall, Value, Name]
 
 	def find_fitting_operator (self, sym, numOperands):
 		# Pass 1: exact numOperands match
+		foundByName = False
 		for op in Operators.values():
 			if op.symbol != sym:
 				continue
 
+			foundByName = True
 			if op.operands == numOperands:
 				return op
 
-		# Pass 2: by symbol
-		for op in Operators.values():
-			if op.symbol == sym:
-				return op
-
-		raise ValueError ('''unknown operator %s!''' % sym)
+		if foundByName:
+			raise ValueError ('''Operator %s used incorrectly''' % sym)
+		else:
+			raise ValueError ('''unknown operator %s!''' % sym)
 
 	def process_operators (self, expr):
 		"""Processes operators"""
@@ -525,7 +574,7 @@
 				args.append (expr[i - 1])
 				argIndices.append (i - 1)
 
-			if i - 1 < len(expr) and self.is_operand (expr[i + 1]):
+			if i + 1 < len(expr) and self.is_operand (expr[i + 1]):
 				args.append (expr[i + 1])
 				argIndices.append (i + 1)
 
@@ -547,6 +596,11 @@
 			else:
 				i += 1
 
+		# Also amend multiplication in sub-expressions
+		for x in expr:
+			if isinstance(x, list):
+				self.amend_multiplication(x)
+
 	def find_priority_operator (self, expr):
 		"""Finds the operator with most priority in the expression"""
 		bestOp = None
@@ -564,6 +618,17 @@
 
 		return (bestOp, bestOpIndex)
 
+	def process_values (self, values):
+		result = []
+		for x in values:
+			if isinstance (x, Name) and not hasattr (x, 'value'):
+				raise ValueError ('%s is unknown' % (x.name))
+			value = x
+			while hasattr (value, 'value'):
+				value = value.value
+			result.append (value)
+		return result
+
 	def evaluate (self, expr, verbose=False):
 		printFunc = print if verbose else lambda x:None
 		printFunc (self.tabs + 'Preprocess: %s' % expr)
@@ -588,7 +653,7 @@
 				self.tabs = self.tabs[:-1]
 
 				printFunc (self.tabs + 'Evaluating function call: %s' % (sym))
-				expr[i] = Value (Functions[sym.funcname]['function'] (*[x.value for x in sym.args]))
+				expr[i] = Value (Functions[sym.funcname]['function'] (*self.process_values(sym.args)))
 				printFunc (self.tabs + '-> %s' % expr[i])
 
 			i += 1
@@ -614,12 +679,19 @@
 			args = [expr[x] for x in argIndices]
 			argIndices = sorted (argIndices, reverse=True)
 			printFunc (self.tabs + 'Processing: (%s, %d) with args %s' % (op, i, args))
-			expr[i] = Value (op.function (*[x.value for x in args]))
+			if op.assign:
+				op.function (*args, calculator=self)
+				expr[i] = AssignmentResult(name=args[0], value=Value(self.variables[args[0].name]))
+			else:
+				expr[i] = Value (op.function (*self.process_values(args)))
 			printFunc (self.tabs + '-> %s' % expr[i])
 
 			for i2 in argIndices:
 				del expr[i2]
 
+			if op.assign and len(expr) != 1:
+				raise ValueError ('expression did not evaluate into an assignment: %s' % expr)
+
 		printFunc (self.tabs + 'Result: %s' % expr[0])
 
 		if len(expr) != 1:
@@ -676,26 +748,35 @@
 
 		if rep == '1':
 			return 'i'
-
-		if rep == '-1':
+		elif rep == '-1':
 			return '-i'
-
-		return rep + 'i'
+		else:
+			return rep + 'i'
 
 	def represent (self, x):
 		"""Returns a string representation of a float or complex number"""
-		if math.fabs (x.imag) > ε:
-			if math.fabs (x.real) > ε:
-				# Complex number
-				return '%s %s %s' % (self.repr_number (x.real),
-					'+' if x.imag >= 0 else '-',
-					self.repr_imaginary (math.fabs (x.imag)))
+		if isinstance (x, Value):
+			x = x.value
+			if isinstance (x, Fraction):
+				return '%s / %s' % (x.numerator, x.denominator)
+			elif math.fabs (x.imag) > ε:
+				if math.fabs (x.real) > ε:
+					# Complex number
+					return '%s %s %s' % (self.repr_number (x.real),
+						'+' if x.imag >= 0 else '-',
+						self.repr_imaginary (math.fabs (x.imag)))
+				else:
+					# Pure imaginary number
+					return self.repr_imaginary (x.imag)
 			else:
-				# Pure imaginary number
-				return self.repr_imaginary (x.imag)
+				# Real number
+				return self.repr_number (x.real)
+		elif isinstance (x, AssignmentResult):
+			return '%s = %s' % (x.name.name, self.represent (x.value))
+		elif isinstance (x, Name) and hasattr(x, 'value'):
+			return self.represent(x.value)
 		else:
-			# Real number
-			return self.repr_number (x.real)
+			return '<%s %s>' % (type(x).__name__, x)
 
 	def calc (self, expr, verbose=None):
 		if verbose is None:
@@ -707,9 +788,10 @@
 		expr = self.tokenize (expr)
 		self.process_parens (expr)
 		self.process_functions (expr)
+		self.process_variables (expr)
 		self.amend_multiplication (expr)
 		self.process_operators (expr)
-		result = self.evaluate (expr, verbose).value
+		result = self.evaluate (expr, verbose)
 		self.previous_result = result
 		return self.represent (result)