add_date.py

# -*- coding: utf-8 -*-
"""
Adiciona períodos a datas
"""

import datetime

DAY = "1"
WEEK = "7"
MONTH = "30"


def add_period(amount, period=DAY, date=None):
    """
    >>> day = datetime.date(2012, 1, 1)
    >>> add_period(1, DAY, day)
    datetime.date(2012, 1, 2)
    >>> add_period(-1, DAY, day)
    datetime.date(2011, 12, 31)
    >>> add_period(0.5, DAY, day)
    datetime.date(2012, 1, 1)
    >>> add_period(-0.5, DAY, day)
    datetime.date(2011, 12, 31)
    >>> add_period(2, DAY, day)
    datetime.date(2012, 1, 3)
    >>> add_period(-2, DAY, day)
    datetime.date(2011, 12, 30)
    >>> add_period(2.5, DAY, day)
    datetime.date(2012, 1, 3)
    >>> add_period(-2.5, DAY, day)
    datetime.date(2011, 12, 29)
    >>> add_period(1, WEEK, day)
    datetime.date(2012, 1, 8)
    >>> add_period(-1, WEEK, day)
    datetime.date(2011, 12, 25)
    >>> add_period(0.5, WEEK, day)
    datetime.date(2012, 1, 4)
    >>> add_period(-0.5, WEEK, day)
    datetime.date(2011, 12, 28)
    >>> add_period(2, WEEK, day)
    datetime.date(2012, 1, 15)
    >>> add_period(-2, WEEK, day)
    datetime.date(2011, 12, 18)
    >>> add_period(2.5, WEEK, day)
    datetime.date(2012, 1, 18)
    >>> add_period(-2.5, WEEK, day)
    datetime.date(2011, 12, 14)
    >>> add_period(1, MONTH, day)
    datetime.date(2012, 2, 1)
    >>> add_period(-1, MONTH, day)
    datetime.date(2011, 12, 1)
    >>> add_period(0.5, MONTH, day)
    datetime.date(2012, 1, 16)
    >>> add_period(-0.5, MONTH, day)
    datetime.date(2011, 12, 16)
    >>> add_period(2, MONTH, day)
    datetime.date(2012, 3, 1)
    >>> add_period(-2, MONTH, day)
    datetime.date(2011, 11, 1)
    >>> add_period(2.5, MONTH, day)
    datetime.date(2012, 3, 16)
    >>> add_period(-2.5, MONTH, day)
    datetime.date(2011, 10, 16)
    >>> add_period(.25, MONTH, day)
    datetime.date(2012, 1, 8)
    >>> add_period(-.25, MONTH, day)
    datetime.date(2011, 12, 24)
    >>> day = datetime.date(2013, 1, 31)
    >>> add_period(1, MONTH, day)
    datetime.date(2013, 2, 28)
    >>> day = datetime.date(2013, 3, 31)
    >>> add_period(-1, MONTH, day)
    datetime.date(2013, 2, 28)
    """
    date = date or datetime.date.today()
    if period == DAY:
        return date + datetime.timedelta(days=amount * 1)
    if period == WEEK:
        return date + datetime.timedelta(days=amount * 7)
    offset = amount > 0 and 32 or 0
    amount = abs(amount)
    while amount:
        step = amount < 1 and amount or 1
        delta = date + datetime.timedelta(days=offset - date.day)
        delta = set_day(delta, date.day) - date
        date += datetime.timedelta(days=delta.days * step)
        amount -= step
    return date


def set_day(date, day):
    """
    Define o dia de um objeto date/datetime e, caso o dia seja inválido, efetua
    algumas tentativas antes de falhar
    >>> set_day(datetime.date(2012, 2, 1), 31)
    datetime.date(2012, 2, 29)
    >>> set_day(datetime.date(2012, 2, 1), -31)
    datetime.date(2012, 2, 29)
    >>> set_day(datetime.date(2012, 2, 1), 0.5)
    datetime.date(2012, 2, 1)
    >>> set_day(datetime.date(2012, 2, 1), "1")
    datetime.date(2012, 2, 1)
    >>> set_day(datetime.date(2012, 2, 1), True)
    datetime.date(2012, 2, 1)
    >>> set_day(datetime.date(2012, 2, 1), False)
    datetime.date(2012, 2, 1)
    >>> set_day(datetime.date(2012, 2, 1), "foo")  # doctest: +ELLIPSIS
    Traceback (most recent call last):
      ...
    TypeError: unsupported operand type(s) for -: 'str' and 'int'
    """
    try:
        day = abs(int(day))
        return date.replace(day=day)
    except ValueError:
        return set_day(date, day - 1)