Adventures in doctesting

parse.py

def operators_and_non_operators(s, operator_list):
    """returns the non-operators
    >>> operators_and_non_operators('2*3+3-4', ['+', '-'])
    (['+', '-'], ['2*3', '3', '4'])
    """
    ops = []
    non_ops = []
    current_non_op = ''
    for c in s:
        if c in operator_list:
            non_ops.append(current_non_op)
            ops.append(c)
            current_non_op = ''
        else:
            current_non_op += c
    non_ops.append(current_non_op)
    return ops, non_ops


def turn_ops_and_parts_into_tree(ops, parts):
    """
    >>> turn_ops_and_parts_into_tree(['*', '-'], ['2', '3', '123'])
    ['-', ['*', '2', '3'], '123']
    """
    tree = parts[0]
    parts = parts[1:]
    while parts:
        tree = [ops.pop(0), tree, parts.pop(0)]
    return tree


def parse(s):
    """
    >>> parse('2*3+4')
    ['+', ['*', '2', '3'], '4']
    >>> parse('2*3+4+5')
    ['+', ['+', ['*', '2', '3'], '4'], '5']
    >>> parse('2*3+4-5+6')
    ['+', ['-', ['+', ['*', '2', '3'], '4'], '5'], '6']
    >>> parse('3*4')
    ['*', '3', '4']
    >>> parse('32')
    '32'
    >>> parse('2+3*4')
    ['+', '2', ['*', '3', '4']]
    """
    if s.isdigit():
        return s
    operators, parts = operators_and_non_operators(s, ['+', '-'])
    parts = [turn_ops_and_parts_into_tree(*operators_and_non_operators(x, ['*', '/'])) for x in parts]
    return turn_ops_and_parts_into_tree(operators, parts)

OPS = {
        '+': lambda x, y: x + y,
        '-': lambda x, y: x - y,
        '*': lambda x, y: x * y,
        '/': lambda x, y: x / y,
        }

def evaluate(tree):
    """

    >>> evaluate(['+', '123', '23'])
    146
    """
    if isinstance(tree, str) and tree.isdigit():
        return int(tree)
    else:
        return OPS[tree[0]](evaluate(tree[1]), evaluate(tree[2]))

def calc(string):
    """The whole thing

    #>>> calc('1+2-4+3')
    2
    #>>> calc('1+2+4-3')
    5
    """
    return evaluate(parse(string))

if __name__ == '__main__':
    import doctest
    doctest.testmod()