Python学习笔记：本月、上月、次月生成

一、需求

因跑模型需要，根据入参，自动设定训练集、测试集、生产集周期。

例如：本月输入“202203”，上月即为“202202”，次月为“202204”。

二、实操

import time, datetime, calendar

# 上月
def last_month(month):
    month = str(month)
    time_tuple = time.strptime(month, "%Y%m")
    year, month, day = time_tuple[:3]
    local_month = datetime.date(year, month, day).replace(day=1)
    last_month = (local_month - datetime.timedelta(days=1)).replace(day=1)
    last_month = last_month.strftime('%Y%m')
    last_month = int(last_month)
    return last_month

# 次月
def next_month(month):
    month = str(month)
    time_tuple = time.strptime(month, "%Y%m")
    year, month, day = time_tuple[:3]
    week, days_num = calendar.monthrange(year, month)
    local_month = datetime.date(year, month, day).replace(day=days_num)
    next_month = (local_month + datetime.timedelta(days=1)).replace(day=1)
    next_month = next_month.strftime('%Y%m')
    next_month = int(next_month)
    return next_month


# 定义取数月份
month = 202203

month_train = month
month_test = last_month(month)
month_predict = next_month(month)
print(r'训练集月份为：', month_train)
print(r'测试集月份为：', month_test)
print(r'预测集月份为：', month_predict)

'''
训练集月份为： 202203
测试集月份为： 202202
预测集月份为： 202204
'''