年、月まで考慮した時間の可算処理

nextperiod.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import calendar
from datetime import datetime
from datetime import timedelta

_FIELD_NAMES = [
    'years',
    'monthes',
    'weeks',
    'days',
    'hours',
    'minutes',
    'seconds',
    'microseconds',
]

MICROSECOND = -1
SECOND = -2
MINUTE = -3
HOUR = -4
DAY = -5
WEEK = -6
MONTH = -7
YEAR = -8


def next_period(dt, field, delta=1):
    if field == YEAR:
        # 年は単純に足すだけ
        return dt.replace(year=dt.year + delta)

    if field == MONTH:
        # 桁があふれうる月 (例えば 13 月とか)
        temporary_month = dt.month - 1 + delta
        # 桁上りを含めた年
        year = dt.year + temporary_month // 12
        # あふれた分を丸めた月
        month = temporary_month % 12 + 1
        # その月の日数
        _, month_last_day = calendar.monthrange(year, month)
        # 元の日付とその月の小さい方を採用する
        day = min(dt.day, month_last_day)
        # 計算結果
        return dt.replace(year=year, month=month, day=day)

    # 週以降は timedelta で自動的に計算できる
    field_name = _FIELD_NAMES[field]
    kwargs = {field_name: delta}
    td = timedelta(**kwargs)

    # 元の日付に差分を足す
    return dt + td


def main():
    # 翌年
    dt = datetime(2016, 1, 1)
    print(next_period(dt, YEAR))  # => 2017-01-01 00:00:00

    # 翌月
    dt = datetime(2016, 1, 1)
    print(next_period(dt, MONTH))  # => 2016-02-01 00:00:00

    # 翌月の同じ日付がない
    dt = datetime(2016, 1, 31)  # 2/31 はないので丸められる
    print(next_period(dt, MONTH))  # => 2016-02-29 00:00:00

    # 翌週 (以降は timedelta を使った実装になるので省略する)
    dt = datetime(2016, 1, 1)
    print(next_period(dt, WEEK))  # => 2016-01-08 00:00:00


if __name__ == '__main__':
    main()