Range function for Decimals and floats and any other objects which supports addition

arange.py

import operator


def arange(start, stop=None, step=None):
    """
    Implement range function not only for integers as Python's builtin
    function, but for Decimals and floats as well.

    Returns generator with arithmetic progession, not list.

        >>> from decimal import Decimal
        >>> list(arange(10))
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
        >>> list(arange(10.))
        [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]
        >>> list(arange(Decimal(10)))
        [Decimal(0), Decimal(1), ..., Decimal(8), Decimal(9)]

        >>> list(arange(1, 11))
        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        >>> list(arange(1., 11.))
        [1., 2., 3., 4., 5., 6., 7., 8., 9., 10.]
        >>> list(arange(Decimal(1), Decimal(11)))
        [Decimal(1), Decimal(2), ..., Decimal(9), Decimal(10)]

        >>> list(arange(0, 30, 5))
        [0, 5, 10, 15, 20, 25, 30]
        >>> list(arange(0., 3., .5))
        [0, .5, 1., 1.5, 2., 2.5, 3.]
        >>> list(arange(Decimal(0), Decimal(3), Decimal('0.5')))
        [Decimal(0), Decimal('0.5'), ..., Decimal('2.5'), Decimal(3)]

        >>> list(arange(0, 10, 3))
        [0, 3, 6, 9]
        >>> list(arange(0., 1., .3))
        [0., .3, .6., .9]
        >>> list(arange(Decimal(0), Decimal(1), Decimal('0.3')))
        [Decimal(0), Decimal('0.3'), Decimal('0.6'), Decimal('0.9')]

        >>> list(arange(0, -10, -1))
        [0, -1, -2, -3, -4, -5, -6, -7, -8, -9]
        >>> list(arange(0, -10., -1.))
        [0, -1., -2., -3., -4., -5., -6., -7., -8., -9.]
        >>> list(arange(Decimal(0), -Decimal(10), -Decimal(1)))
        [Decimal(0), -Decimal(1), ..., -Decimal(8), -Decimal(9)]

        >>> list(arange(0))
        []
        >>> list(arange(0.))
        []
        >>> list(arange(Decimal()))
        []

        >>> list(arange(1, 0))
        []
        >>> list(arange(1., 0.))
        []
        >>> list(arange(Decimal(1), Decimal(0)))
        []

    """
    klass = type(start)
    lt_func = operator.lt

    stop = start if stop is None else stop
    start = klass(0) if start == stop else start
    step = klass(1 if step is None else step)

    assert isinstance(stop, klass), (
        'Start and stop limits have different types, {0!r} != {1!r}.'.
        format(type(start).__name__, type(stop).__name__)
    )
    assert step, "Step shouldn't be a zero: {0!r}.".format(step)

    if start < stop and step < 0 or start > stop and step > 0:
        raise StopIteration
    elif start > stop and step < 0:
        lt_func = operator.gt

    while lt_func(start, stop):
        yield start
        start += step

tests.py

from decimal import Decimal
from unittest import TestCase

from arange import arange


class TestArange(TestCase):

    def test_arange(self):
        self.assertRaises(AssertionError, list, arange(1, 2.))
        self.assertRaises(AssertionError, list, arange(1, 2, 0))

        self.assertEqual(list(arange(10)), range(10))
        self.assertEqual(list(arange(10.)), map(float, range(10)))
        self.assertEqual(list(arange(Decimal(10))), map(Decimal, range(10)))

        self.assertEqual(list(arange(1, 11)), range(1, 11))
        self.assertEqual(list(arange(1., 11.)), map(float, range(1, 11)))
        self.assertEqual(list(arange(Decimal(1), Decimal(11))),
                         map(Decimal, range(1, 11)))

        self.assertEqual(list(arange(0, 30, 5)), range(0, 30, 5))
        self.assertEqual(list(arange(0., 3., .5)),
                         map(lambda item: float(item) / 10., range(0, 30, 5)))
        self.assertEqual(list(arange(Decimal(), Decimal(3), Decimal('.5'))),
                         map(lambda item: Decimal(item) / 10, range(0, 30, 5)))

        self.assertEqual(list(arange(0, 10, 3)), range(0, 10, 3))

        int_range = range(0, 10, 3)

        for i, item in enumerate(arange(0., 1., .3)):
            self.assertAlmostEqual(item, float(int_range[i]) / 10.)

        self.assertEqual(list(arange(Decimal(), Decimal(1), Decimal('.3'))),
                         map(lambda item: Decimal(item) / 10, range(0, 10, 3)))

        self.assertEqual(list(arange(0, -10, -1)), range(0, -10, -1))
        self.assertEqual(list(arange(0., -10., -1)),
                         map(float, range(0, -10, -1)))
        self.assertEqual(list(arange(Decimal(), -Decimal(10), -1)),
                         map(Decimal, range(0, -10, -1)))

        self.assertEqual(list(arange(0)), [])
        self.assertEqual(list(arange(0.)), [])
        self.assertEqual(list(arange(Decimal())), [])

        self.assertEqual(list(arange(1, 0)), [])
        self.assertEqual(list(arange(1., 0.)), [])
        self.assertEqual(list(arange(Decimal(1), Decimal())), [])

        self.assertEqual(list(arange(1, 2, -1)), [])
        self.assertEqual(list(arange(1, -1)), [])