Equation parser and solver

Question

I wrote this equation parser and solver, I feel like it is well documented and tested:

import doctest def swap_sign(num): """ Swaps the sign of a number-representing string. >>> swap_sign('9') '-9' >>> swap_sign('-4x') '4x' """ if num.startswith('-'): return ''.join(num[1:]) elif num.startswith('+'): return '-'+''.join(num[1:]) else: return '-'+num def add_one_before_x(expr): """ It is a common usage to omit the one before the x saying 'x' when you really mean '1x'. >>> add_one_before_x('x +2x = 9 -x^2') '1x +2x = 9 -1x^2' """ if expr.startswith('x'): expr = '1'+expr return expr.replace(' x',' 1x').replace(' -x',' -1x').replace(' +x',' +1x') def move_operators(expr): """ Sticks the operators the next term for easier further parsing. >>> move_operators('3 + 9x = 3 - 3x^2 + 4') '3 +9x = 3 -3x^2 +4' """ return expr.replace('+ ','+').replace('- ','-').replace('--','-') def move_to_the_same_side(expr): """ Moves all the numbers and x-s to one side, changing sign when needed. >>> move_to_the_same_side("3x -4x^2 +5 -2 = 2x") ['3x', '-4x^2', '+5', '-2', '-2x'] >>> move_to_the_same_side("9x = 1x^2 -4 +2 -1 +7x") ['9x', '-1x^2', '4', '-2', '1', '-7x'] """ if ' = ' not in expr: raise ValueError("There is no second term, remember spaces around the equal sign.") left,right = expr.split(' = ') return left.split(' ') + [swap_sign(token) for token in right.split(' ')] def to_normal_form(expr): """ Performs further parsing on all the coefficients on one side and return the coefficients (a,b,c). >>> to_normal_form(['3x', '-4x^2', '+5', '-2', '-2x']) (-4, 1, 3) """ bare_nums = [i for i in expr if 'x' not in i] xes = [i for i in expr if 'x' in i and '^2' not in i] two_pow_xes = [i for i in expr if 'x' in i and '^2' in i] final_num = sum([int(n) for n in bare_nums]) final_x = sum([int(x.replace('x','')) for x in xes]) final_pow_x = sum([int(x.replace('x^2','')) for x in two_pow_xes]) return final_pow_x, final_x, final_num def first_grade_solve(coeffs): """ Solves the first grade equations using the trivial first grade formula. Also solves equations of grade 0. >>> first_grade_solve((0,0,3)) # 0x + 3 = 0 'Never' >>> first_grade_solve((0,2,3)) # 2x + 3 = 0 (-1.5,) """ _,a,b = coeffs if a == 0: return "Every time" if b == 0 else "Never" return ((-b) / a,) def second_grade_solve(coeffs): """ Solves second grade equations using the well known second grade formula. >>> second_grade_solve((1,2,1)) # 1x^2 + 2x + 1 = 0 (-1.0, -1.0) >>> second_grade_solve((1,5,6)) # 1x^2 + 5x + 6 = 0 (-2.0, -3.0) """ a,b,c = coeffs delta = b**2 - 4*a*c return (((-b)+delta**.5)/(2*a), ((-b)-delta**.5)/(2*a)) def solve(coefficients): """ Dispatches solving to the correct method or aborts if the equation grade is too high. >>> solve((1,5,6)) (-2.0, -3.0) >>> solve((0,2,4)) (-2.0,) """ if coefficients[0] == 0: return first_grade_solve(coefficients) elif len(coefficients) == 3: return second_grade_solve(coefficients) raise NotImplementedError("Only 0th, 1st and 2nd grade equations can be solved") def parse_and_solve(expr): """ Connects the other functions to provide full equation solving. >>> parse_and_solve("2x - 4 = 0") (2.0,) >>> parse_and_solve("1 = 2") 'Never' >>> parse_and_solve("x^2 + 2x = -1") (-1.0, -1.0) """ simpler = add_one_before_x(move_operators(expr)) same_sided = move_to_the_same_side(simpler) normal_form = to_normal_form(same_sided) return solve(normal_form) def interface(): print("Welcome to the symbolic equation solver.") print("Please enter equations in the form:") print("3x - 4x^2 + 5 - 2 = 2x + x^2\n") while True: expr = input("> ") try: result = parse_and_solve(expr) if result in ("Every time","Never"): print(result) else: print('\n'.join(['x = '+str(x) for x in result])) except Exception as E: print("Invalid Expression:", E) if __name__ == "__main__": doctest.testmod() interface() 

Answer 1

1. This is well documented code, clearly written and sensibly factored into functions.

2. The docstrings need a bit of work. It's best to imagine reading a docstring from the point of view of a user who wants to know how to call a function and so has just typed help(function) in the Python interpreter. But if they see something like this:

   >>> help(solve) Help on function solve in module __main__: solve(coefficients) Dispatches solving to the correct method or aborts if the equation grade is too high. 

   then they are none the wiser. What do I pass as coefficients? What does it return? What does it mean by "correct method"? The user was hoping for something more like this:

   solve(coefficients) Given the sequence of coefficients of x^0, x^1, ..., in a polynomial, returns a tuple of roots of that polynomial. Raises NotImplementedError for polynomials of degree greater than 2. 

3. I think degree is the more usual term, not grade.

4. first_grade_solve sometimes returns a length-1 tuple of solutions, and sometimes returns a string. It's not a good idea to write functions that return different kinds of result in exceptional situations, because it's too easy for the caller to forget to handle the exceptional case and write something like:

   x = first_grade_solve(coeffs)[0] 

   which will mistakenly set x to the string 'N' if the equation has no solution. It's better to deal with an exceptional situation by raising an exception.

5. Running all the doctests each time the program is executed is a mistake. If the user just wants to run the program then they probably don't want to wait for the tests to finish running first. On the other hand, if they just want to test the program they probably don't want to have to deal with the interactive interface. So just run the program; a user who wants to run the doctests can easily do so using the -m command-line option:

   python -m doctest program.py 

6. The parsing relies on repeated string substitution, which is hard to make reliable:

   >>> parse_and_solve('x = 1') Traceback (most recent call last): ... final_num = sum([int(n) for n in bare_nums]) ValueError: invalid literal for int() with base 10: '' 

   and is prone to false positives:

   >>> parse_and_solve('xx = 1') (1.0,) >>> parse_and_solve('x2x = 12') (1.0,) 

   It's usually more reliable and robust to use a real parser.

Here's how you might implement an equation parser using PyParsing:

from collections import Counter from pyparsing import Group, Literal, Optional, Regex, Suppress, ZeroOrMore def equation_grammar(): """Return a PyParsing grammar for a polynomial equation.""" plus = Literal('+').setParseAction(lambda:1) minus = Literal('-').setParseAction(lambda:-1) sign = plus | minus number = Regex('[0-9]+').setParseAction(lambda t:int(t[0])) exponent = Optional(Suppress('^') + number, default=1)('exponent') power = 'x' + exponent.setParseAction(lambda t:t[0]) atom = number('coeff') + Optional(power) | power monomial = Optional(sign('sign')) + atom polynomial = Group(monomial) + ZeroOrMore(Group(sign('joiner') + monomial)) return polynomial('left') + '=' + polynomial('right') _grammar = equation_grammar() def parse(s): """Parse the string s as a polynomial equation in x; gather terms to express it as a power series in x equated to zero; and return a list of the coefficients. >>> parse('x = 1') [1, -1] >>> parse('7x^7 + -5x^5 + 3x^3 + -1 = 3x^6 - x^4 + 7') [7, -3, -5, 1, 3, 0, 0, -8] """ equation = _grammar.parseString(s) series = Counter() for side, terms in ((+1, equation.left), (-1, equation.right)): for t in terms: n = side * t.get('joiner', 1) * t.get('sign', 1) * t.get('coeff', 1) series[t.get('exponent', 0)] += n degree = max(e for e, c in series.items() if c) return [series[e] for e in range(degree, -1, -1)] 

I've chosen the output format of parse so that it can be fed directly into numpy.roots:

>>> from numpy import roots >>> roots(parse('7x = 21')) array([ 3.]) >>> roots(parse('x^2 + 3x = -2')) array([-2., -1.]) >>> roots(parse('2x^3 + 22x = 12x^2 + 12')) array([ 3., 2., 1.]) >>> roots(parse('x^5 = 1')) array([-0.80901699+0.58778525j, -0.80901699-0.58778525j, 0.30901699+0.95105652j, 0.30901699-0.95105652j, 1.00000000+0.j ])