datetimexrange.py

# -*- coding: utf-8 -*-
from copy import copy
from datetime import timedelta


ONE_DAY = timedelta(days=1)


class datetimerange(object):
    """Just like ``xrange``, but working on ``datetime``."""

    weekend_days = (5, 6)

    def __init__(self, from_, to, step=ONE_DAY, exclude_weekend=False):
        self.from_ = from_
        self.current = from_
        self.limit = to
        self.step = step
        self.exclude_weekend = exclude_weekend

    def __iter__(self):
        return self.__class__(
            self.from_,
            self.limit,
            self.step,
            exclude_weekend=self.exclude_weekend,
        )

    def next(self):
        if self.current >= self.limit:
            raise StopIteration()
        current = copy(self.current)
        _next = current + self.step
        if self.exclude_weekend:
            while current.weekday() in self.weekend_days:
                current += self.step
                _next += self.step
        self.current = _next
        return (current, _next)