交易日工具类

utils.py

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

"""utils.py: 交易日工具

__email__ = "[email protected]"
__date__ = "2021-12-21"
__version__ = "1.0.0"
"""

from datetime import datetime, timedelta

fmt = '%Y%m%d'
holidays = [
    '20220101', '20220102', '20220103',  # 元旦
    '20220131', '20220201', '20220202', '20220203', '20220204', '20220205', '20220206',  # 春节
    '20220403', '20220404', '20220405',  # 清明节
    '20220430', '20220501', '20220502', '20220503', '20220504',  # 劳动节
    '20220603', '20220604', '20220605',  # 端午节
    '20220910', '20220911', '20220912',  # 中秋节
    '20221001', '20221002', '20221003', '20221004', '20221005', '20221006', '20221007'  # 国庆节
]


def get_trade_day(str_date, delta: int):
    """获取指定日期相差 delta 天的交易日期

    如果 delta 为 0，则判断给定日期是否交易日，如果不是交易日，返回None;是交易日，则返回自身。

    :param str_date: 指定的日期字符串，格式 YYYYMMDD
    :param delta: 相差的天数。如果需要历史日期，填写为负整数；如果是将来的日期，为正整数
    :return: 返回具体的日期字符串，格式 YYYYMMDD
    """
    date = datetime.strptime(str_date, fmt)
    direction = 1 if delta > 0 else -1

    while delta != 0:
        date = date + timedelta(days=1 * direction)
        wd = date.weekday()
        dt = date.strftime(fmt)
        # print(f'weekday={wd}, dt={dt}')
        if wd not in [5, 6] and dt not in holidays:
            delta += -1 * direction
            continue
    else:
        wd = date.weekday()
        dt = date.strftime(fmt)
        # print(f'weekday={wd}, dt={dt}')
        if wd in [5, 6] or dt in holidays:
            return None
    return date.strftime(fmt)


def get_pre_trade_day(str_date):
    """获取 str_date 日期的上一交易日期

    :param str_date: 指定的日期字符串，格式 YYYYMMDD
    :return: 返回具体的日期字符串，格式 YYYYMMDD
    """
    return get_trade_day(str_date, -1)


def get_next_trade_day(str_date):
    """获取 str_date 日期的下一交易日期

    :param str_date: 指定的日期字符串，格式 YYYYMMDD
    :return: 返回具体的日期字符串，格式 YYYYMMDD
    """
    return get_trade_day(str_date, 1)


def is_trade_day(str_date):
    """获取 str_date 特定日期

    :param str_date: 指定的日期字符串，格式 YYYYMMDD
    :return: 是否交易日
    """
    return get_trade_day(str_date, 0) == str_date


def get_delta_days(begin_date, end_date):
    """获取两个日期之间相差天数

    :param begin_date: 开始日期
    :param end_date: 结束日期
    :return: 天数的绝对值
    """
    d1 = datetime.strptime(begin_date, fmt)
    d2 = datetime.strptime(end_date, fmt)
    delta = (d1 - d2).days
    return delta if delta > 0 else delta * -1

utils_test.py

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

from utils import get_pre_trade_day, get_next_trade_day, is_trade_day


def test_get_pre_trade_day():
    a = get_pre_trade_day('20220104')
    assert a == '20211231'
    a = get_pre_trade_day('20211220')
    assert a == '20211217'
    a = get_pre_trade_day('20211217')
    assert a == '20211216'
    a = get_pre_trade_day('20221007')
    assert a == '20220930'


def test_is_trade_day():
    a = is_trade_day('20211231')
    assert a is True
    a = is_trade_day('20221007')
    assert a is False


def test_get_next_trade_day():
    a = get_next_trade_day('20221217')
    assert a == '20221219'
    a = get_next_trade_day('20220930')
    assert a == '20221010'
    a = get_next_trade_day('20220908')
    assert a == '20220909'