用python進行月份加減的函數

import math


def add_month(datamonth, num):
    """
    月份加減函數,返回字符串類型
    :param datamonth: 時間(201501)
    :param num: 要加(減)的月份數量
    :return: 時間(str)
    """
    months = ['01', '02', '03', '04', '05', '06', '07', '08', '09', '10', '11', '12']
    datamonth = int(datamonth)
    num = int(num)
    year = datamonth // 100
    new_list = []
    s = math.ceil(abs(num) / 12)
    for i in range(int(-s), s + 1):
        new_list += [str(year + i) + x for x in months]
    new_list = [int(x) for x in new_list]
    return str(new_list[new_list.index(datamonth) + num])


if __name__ == '__main__':
    a = add_month('201501', 5)
    print(a)