用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)