Compute the sum of a range without iterating it

sumrange.py

"""
Helper to compute the sum of a range without iterating it.

Equivalent to sum(range(start, stop, step)).
"""


def sumrange(start, stop=None, step=1):
    # Handle args
    if stop is None:
        start, stop = 0, start
    # Compute the length of the range
    sign = (step > 0) - (step < 0)
    length = max(0, 1 + (stop - start - sign) // step)
    # Sum of the first n integers starting at 0
    partial_sum = length * (length - 1) // 2
    # Apply offset and step
    return start * length + step * partial_sum


def test(n=100):
    for stop in range(-n, n):
        expected = sum(range(stop))
        assert sumrange(stop) == expected

    for start in range(-n, n):
        for stop in range(-n, n):
            expected = sum(range(start, stop))
            assert sumrange(start, stop) == expected

    for start in range(-n, n):
        for stop in range(-n, n):
            for step in range(1, n):
                expected = sum(range(start, stop, step))
                assert sumrange(start, stop, step) == expected

    for start in range(-n, n):
        for stop in range(-n, n):
            for step in range(-n, 0):
                expected = sum(range(start, stop, step))
                assert sumrange(start, stop, step) == expected