Python求前后第n个月函数

数据分析常用到：

# 求往后第 n 个月
# monthstr 格式：YYYY-MM
#
# print next_month("2023-05", 1)
# 2023-06
# print next_month("2023-05", 6)
# 2023-11
# print next_month("2023-05", 8)
# 2024-01
def next_month(monthstr, n):
    yearnum = int(monthstr[0:4])
    monthnum = int(monthstr[5:7])
    date = datetime.datetime.strptime(monthstr + "-01", '%Y-%m-%d')    
    for i in range(0, n):
        if monthnum == 12:
            yearnum += 1
            monthnum = 1
        else:
            monthnum += 1
    return "{}-{:02}".format(yearnum, monthnum)

# 求往前第 n 个月
# monthstr 格式：YYYY-MM
#
# print prev_month("2023-05", 1)
# 2023-04
# print prev_month("2023-05", 3)
# 2023-02
# print prev_month("2023-05", 5)
# 2022-12
def prev_month(monthstr, n):
    yearnum = int(monthstr[0:4])
    monthnum = int(monthstr[5:7])
    date = datetime.datetime.strptime(monthstr + "-01", '%Y-%m-%d')    
    for i in range(0, n):
        if monthnum == 1:
            yearnum -= 1
            monthnum = 12
        else:
            monthnum -= 1
    return "{}-{:02}".format(yearnum, monthnum)