Source

flowrate / flowrate / cmpdates.py

import datetime


class month(object):
    """A given month in a given year. Like datetime.date without a .day field."""

    __slots__ = ('year', 'month')

    def __init__(self, year, month):
        self.year = year
        self.month = month

    @classmethod
    def fromdate(cls, dt):
        return cls(dt.year, dt.month)

    def todate(self, day=1):
        return datetime.date(self.year, self.month, day)

    def __add__(self, other):
        # Add an integer number of months to self.
        y, m = divmod(self.month + other, 12)
        return month(self.year + y, m)

    def __sub__(self, other):
        if isinstance(other, month):
            # Subtract a month() instance to return an integer diff of months
            y, m = divmod(self.month - other.month, 12)
            return (y * 12) + m
        else:
            # Subtract an integer number of months from self.
            y, m = divmod(self.month - other, 12)
            return month(self.year + y, m)

    def __cmp__(self, other):
        return cmp((self.year, self.month), (other.year, other.month))


class week(object):
    """A given week in a given year. Like datetime.date with a .week field."""

    __slots__ = ('year', 'week')

    def __init__(self, year, week):
        self.year = year
        self.week = week

    @classmethod
    def fromdate(cls, dt):
        y, w, d =  dt.isocalendar()
        return cls(y, w)

    def todate(self, day=1):
        jan1 = datetime.date(self.year, 1, 1)
        day1ofweek1 = jan1 + datetime.timedelta(days=8 - jan1.isocalendar()[2])
        return (day1ofweek1 +
                datetime.timedelta(
                    days=((self.week - 1) * 7) + (day - 1)
                ))

    def __add__(self, other):
        # Add an integer number of weeks to self.
        f = self.todate() + datetime.timedelta(days=other * 7)
        y, w, d = f.isocalendar()
        return week(y, w)

    def __sub__(self, other):
        if isinstance(other, week):
            # Subtract a week() instance to return an integer diff of weeks
            return ((self.todate() - other.todate()).days / 7)
        else:
            # Subtract an integer number of weeks from self.
            f = self.todate() - datetime.timedelta(days=other * 7)
            y, w, d = f.isocalendar()
            return week(y, w)

    def __cmp__(self, other):
        return cmp((self.year, self.week), (other.year, other.week))


# A map from 'units' (time periods) to coercion_funcs.
# The coercion_func is used to coerce two dates to a common representation
# at the granularity for the units. For example, if 'units' is 'months',
# we want date(2012, 5, 1) to compare equal with date(2012, 5, 15).
units = {
    'days': lambda d: d,
    'weeks': week.fromdate,
    'months': month.fromdate,
    'years': lambda d: d.year,
    }
# The day_func is used to determine if a date matches our 'days' value,
# which, depending on the 'units', can mean "day of the month",
# "day of the week", or "day of the year".
days = {
    'days': lambda d: 1,
    'weeks': lambda d: d.weekday(),
    'months': lambda d: d.day,
    'years': lambda d: d.timetuple().tm_yday,
    }

# A map from units to coercion funcs which take dates and return strings.
# These can then be compared to each and also used as dict keys for fast
# lookups (but cannot be added and subtracted).
strkeys = {
    'days': lambda d: d.isoformat(),
    'weeks': lambda d: '%04d,%02d' % d.isocalendar()[:2],
    'months': lambda d: '%04d-%02d' % (d.year, d.month),
    'years': lambda d: '%04d' % d.year,
    }