Toolbox: Get date n weeks from another date in Python

The function uses only basic Python modules (datetime and timedelta). If you want to:

• trim your weekly time series into training and test set,
• or create or check the date n weeks ahead or back,

then this function is for you.

from datetime import datetime, timedelta


def get_date_n_weeks(fixed_date, weeks_diff, dateformat=None, direction='back'):
    """
    Function returns specific date n weeks ahead or back.

    INPUT:

    :param fixed_date: (str or datetime) date from which other date will be calculated.
    :param weeks_diff: number of weeks back or ahead from the fixed date.
    :param dateformat: (str, optional) if given fixed_date has type string then parameter must be
                       set to the date format.
    :param direction: (str) 'back' or 'ahead'.

    OUTPUT:

    :returns: (datetime) other date n weeks ahead or back
    """

    if isinstance(fixed_date, str):
        date = datetime.strptime(fixed_date, dateformat)
    else:
        date = fixed_date

    d = timedelta(days=weeks_diff * 7)  # 7 days in one week

    if direction == 'back':
        new_date = date - d
    elif direction == 'ahead':
        new_date = date + d
    else:
        raise TypeError('Wrong parameter: direction parameter could be only "back" or "ahead".')

    return new_date