daterange.py

# -*- coding: utf-8 -*-

"""
Example Usage
=============

>>> import datetime
>>> start = datetime.date(2009, 6, 21)

>>> g1 = daterange(start)
>>> g1.next()
datetime.date(2009, 6, 21)
>>> g1.next()
datetime.date(2009, 6, 22)
>>> g1.next()
datetime.date(2009, 6, 23)
>>> g1.next()
datetime.date(2009, 6, 24)
>>> g1.next()
datetime.date(2009, 6, 25)
>>> g1.next()
datetime.date(2009, 6, 26)

>>> g2 = daterange(start, to=datetime.date(2009, 6, 25))
>>> g2.next()
datetime.date(2009, 6, 21)
>>> g2.next()
datetime.date(2009, 6, 22)
>>> g2.next()
datetime.date(2009, 6, 23)
>>> g2.next()
datetime.date(2009, 6, 24)
>>> g2.next()
datetime.date(2009, 6, 25)
>>> g2.next()
Traceback (most recent call last):
...
StopIteration

>>> g3 = daterange(start, step=datetime.timedelta(days=2))
>>> g3.next()
datetime.date(2009, 6, 21)
>>> g3.next()
datetime.date(2009, 6, 23)
>>> g3.next()
datetime.date(2009, 6, 25)
>>> g3.next()
datetime.date(2009, 6, 27)

>>> g4 = daterange(start, to=datetime.date(2009, 6, 25),
...     step=datetime.timedelta(days=2))
>>> g4.next()
datetime.date(2009, 6, 21)
>>> g4.next()
datetime.date(2009, 6, 23)
>>> g4.next()
datetime.date(2009, 6, 25)
>>> g4.next()
Traceback (most recent call last):
...
StopIteration

"""

import datetime
import re


def daterange(date, to=None, step=datetime.timedelta(days=1)):
    """
    Similar to the built-in ``xrange()``, only for datetime objects.
    
    If called with just a ``datetime`` object, it will keep yielding values
    forever, starting with that date/time and counting in steps of 1 day.
    
    If the ``to_date`` keyword is provided, it will count up to and including
    that date/time (again, in steps of 1 day by default).
    
    If the ``step`` keyword is provided, this will be used as the step size
    instead of the default of 1 day. It should be either an instance of
    ``datetime.timedelta``, an integer, or a string representing an integer.
    If it is either of the latter two, it will be interpreted as a number of
    days.
    
    Note that, due to the similar interfaces of both objects, this function
    will accept both ``datetime.datetime`` and ``datetime.date`` objects.
    """    
    if to is None:
        condition = lambda d: True
    else:
        condition = lambda d: (d <= to)
    
    if (isinstance(step, (int, long)) or re.match(r'^([\d]+)$', str(step))):
        # By default, integers are interpreted in days. For more granular
        # steps, use a `datetime.timedelta()` instance.
        step = datetime.timedelta(days=int(step))
    
    if not isinstance(step, datetime.timedelta):
        raise TypeError('Invalid step value: %r' % (step_size,))
    
    # The main generation loop.
    while condition(date):
        yield date
        date += step


if __name__ == '__main__':
    import doctest
    doctest.testmod()