calculator.py

"""
String Calculator

    1- The program can take 0, 1 or 2 numbers and will
    return their sum, (for an empty string it will return 0)
    for example "" or "1" or "1,2" -> 0, 1, 3

    2- Allow the Add method to handle an unknown amount
    of numbers

    3- Allow the Add method to handle new lines between
    numbers(instead of commas). The following input
    is ok: "1\n2,3" (will equal 6)

    4- Support different delimiters. To change a delimiter
    the beginning of the string will contain a separate
    line that looks like this: "//[delimiter]\n[numbers...]"
    for example "//;\n1;2" should return 3 where the
    default delimiter is ";"
"""
import itertools

def add(text):
    if text == "":
        return 0

    numbers = _get_numbers(text)

    return sum(numbers)

def _get_numbers(text):
    delimiter = ","
    lines = text.splitlines()

    if lines[0].startswith("//"):
        delimiter = _get_delimiter(lines[0])
        lines = lines[1:]

    return map(int, _split(lines, delimiter))

def _get_delimiter(line):
    return line[2:]

def _split(lines, delimiter):
    return itertools.chain.from_iterable(
        line.split(delimiter) for line in lines)

test_calculator.py

from calculator import add

def test_should_return_zero_on_empty_string():
    assert add("") == 0

def test_should_return_self_on_number():
    assert add("1") == 1
    assert add("2") == 2

def test_should_return_sum_on_two_numbers():
    assert add("1,2") == 3
    assert add("3,3") == 6

def test_should_return_sum_on_many_numbers():
    assert add("1,2,3") == 6
    assert add("1,2,3,4") == 10

def test_should_accept_linebreak_as_separator():
    assert add("1\n2") == 3
    assert add("1\n2,3") == 6

def test_should_accept_custom_delimiters():
    assert add("//;\n1;2") == 3
    assert add("//;\n1;2\n3") == 6