воскресенье, 4 мая 2014 г.

Range для не целых чисел

Почему-то никогда не думал, что rangexrange) понимают только целые числа (в Python 3 ситуация такая же) и функцией нельзя воспользоваться, как например:

range(0., 5., .5)
range(Decimal('0'), Decimal('10'), Decimal('1.5'))

Поэтому пришлось сделать свою замену:

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.
    """
    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

В отличии от range в Python 2 arange отдает генератор, а не список, что вообщем-то удобней и правильней, ну и функция всегда вернет элементы такого же типа, как были переданы в start и stop аргументы, даже если шаг - это целое число или он не указан вовсе.

Как всегда код оформлен как Gist на GitHub'е, там еще доктесты и юниттесты.

blog comments powered by Disqus