date_range.py

#!/usr/bin/env python

__doc__ = """
This module defines a simple date_range class, this is an iterator over
a range of dates starting at one day and ending at another
"""

from datetime import datetime, date, timedelta

# Our date_range classm this is what it's all about
class date_range(object):
    __doc__ = """This class defines the iterator
Constructor parameters:
 since = Starting date (datetime or date)
 until = Ending date (datetime or date)
 step = Step size (timedelta), optional, defaults to 1 day
"""

    def __init__(self, since, until, step=timedelta(days=1)):
        assert isinstance(since, (datetime, date))
        assert isinstance(until, (datetime, date))
        assert isinstance(step, timedelta)
        # Infinite loop risk if the step value is zero
        assert (step != timedelta(0))

        self.since = since
        self.until = until
        self.step = step
        # Allow for negative step value
        if self.step > timedelta(0):
            # Make sure since is before until for positive step
            if self.since > self.until:
                temp = self.until
                self.until = self.since
                self.since = temp
        else:
            # Make sure since is after until for negative step
            if self.since < self.until:
                temp = self.until
                self.until = self.since
                self.since = temp
        self.current = self.since - self.step

    def __iter__(self):
        return self

    def next(self):
        # Allow for negative step value
        if (self.step > timedelta(0)):
            if self.current >= self.until:
                raise StopIteration()
        else:
            if self.current <= self.until:
                raise StopIteration()

        self.current += self.step

        # Allow for negative step value
        if (self.step > timedelta(0)):
            if self.current > self.until:
                return self.until
        else:
            if self.current < self.until:
                return self.until

        return self.current


def test():
    d_range = list(date_range(date(2000, 1, 1), date(2000, 1, 3)))
    e_range = [date(2000, 1, 1), date(2000, 1, 2), date(2000, 1, 3)]
    assert d_range == e_range, d_range

    d_range = list(date_range(
        date(2000, 1, 1),
        date(2000, 2, 3),
        timedelta(days=7)
    ))
    e_range = [
        date(2000, 1, 1),
        date(2000, 1, 8),
        date(2000, 1, 15),
        date(2000, 1, 22),
        date(2000, 1, 29),
        date(2000, 2, 3),
    ]
    assert d_range == e_range, d_range
    # This test does nasty stuff, negative step value, since and until in the wrong order
    # The class should just silently fix it
    d_range = list(date_range(
        date(2013,1,1),
        date(2013,1,3),
        timedelta(days=-1)
    ))
    e_range = [
        date(2013,1,3),
        date(2013,1,2),
        date(2013,1,1)
    ]
    assert d_range == e_range, d_range


if __name__ == '__main__':
    test()