量化交易之One Piece篇 - onepiece_rsh - trading_time.json处理添加日内模式

from tqz_extern.json_operator import TQZJsonOperator
import os
from datetime import datetime, timedelta

from tqz_extern.tqz_constant import TQZTradingTimeType


class MakeTradingTime:

    DAY: str = 'day'
    NIGHT: str = 'night'

    day_close_times: list = ["1500", "1515"]
    night_close_times: list = ["2300", "0100", "0230"]

    @classmethod
    def reset_trading_time_json(cls, source_trading_time_json: str, target_trading_time_json: str, offset_close_minutes: int = 0, trading_time_type: TQZTradingTimeType = TQZTradingTimeType.DAY_TYPE):
        assert os.path.exists(path=source_trading_time_json), f'Bad path: {source_trading_time_json} not exist.'
        trading_time_content = TQZJsonOperator.tqz_load_jsonfile(jsonfile=source_trading_time_json)

        target_content: {str, dict} = {}
        if trading_time_type in [TQZTradingTimeType.DAY_SESSION_TYPE, TQZTradingTimeType.ALL_SESSION_TYPE]:
            target_content = cls.__get_session_trading_time_content(
                source_trading_time_content=trading_time_content,
                offset_close_minutes=offset_close_minutes,
                trading_time_type=trading_time_type
            )
        elif trading_time_type in [TQZTradingTimeType.DAY_TYPE]:
            target_content = cls.__get_day_trading_time_content(
                source_trading_time_content=trading_time_content,
                offset_close_minutes=offset_close_minutes,
                trading_time_type=trading_time_type
            )

        TQZJsonOperator.tqz_write_jsonfile(content=target_content, target_jsonfile=target_trading_time_json)


    @classmethod
    def __get_day_trading_time_content(cls, source_trading_time_content, offset_close_minutes: int = 0, trading_time_type: TQZTradingTimeType = TQZTradingTimeType.DAY_SESSION_TYPE) -> {str, dict}:
        assert trading_time_type in [TQZTradingTimeType.DAY_TYPE], f'Bad trading_time_type: {trading_time_type}'

        target_content: {str, dict} = {}
        for instrument, data in source_trading_time_content.items():
            target_content[instrument] = {cls.DAY: [], cls.NIGHT: []}
            session_item = []
            for session in data[cls.DAY]:
                for session_time in session:
                    session_start_end = f'{session_time.split(":")[0]}{session_time.split(":")[1]}'
                    if session_start_end in cls.day_close_times:
                        new_time_obj = datetime.strptime(session_start_end, "%H%M") - timedelta(minutes=offset_close_minutes)

                        session_item.append(int(new_time_obj.strftime("%H%M")))
                    else:
                        session_item.append(int(session_start_end))

            target_content[instrument][cls.DAY].append([session_item[0], session_item[-1]])

        return target_content


    @classmethod
    def __get_session_trading_time_content(cls, source_trading_time_content, offset_close_minutes: int = 0, trading_time_type: TQZTradingTimeType = TQZTradingTimeType.DAY_SESSION_TYPE) -> {str, dict}:
        assert trading_time_type in [TQZTradingTimeType.DAY_SESSION_TYPE, TQZTradingTimeType.ALL_SESSION_TYPE], f'Bad trading_time_type: {trading_time_type}'

        target_content: {str, dict} = {}
        for instrument, data in source_trading_time_content.items():
            target_content[instrument] = {cls.DAY: [], cls.NIGHT: []}
            for session in data[cls.DAY]:
                session_item = []
                assert 2 == len(session), f'Bad session size: {session}'
                session_start = f'{session[0].split(":")[0]}{session[0].split(":")[1]}'
                session_item.append(int(session_start))

                session_end = f'{session[1].split(":")[0]}{session[1].split(":")[1]}'
                new_time_obj = datetime.strptime(session_end, "%H%M") - timedelta(minutes=offset_close_minutes)
                session_item.append(int(new_time_obj.strftime("%H%M")))

                target_content[instrument][cls.DAY].append(session_item)

            if trading_time_type == TQZTradingTimeType.ALL_SESSION_TYPE:
                for session in data[cls.NIGHT]:
                    session_item = []
                    assert 2 == len(session), f'Bad session size: {session}'
                    session_start = f'{session[0].split(":")[0]}{session[0].split(":")[1]}'
                    session_item.append(int(session_start))

                    session_end = f'{session[1].split(":")[0]}{session[1].split(":")[1]}'
                    new_time_obj = datetime.strptime(session_end, "%H%M") - timedelta(minutes=offset_close_minutes)
                    session_item.append(int(new_time_obj.strftime("%H%M")))

                    target_content[instrument][cls.NIGHT].append(session_item)

        return target_content


if __name__ == '__main__':
    MakeTradingTime.reset_trading_time_json(
        source_trading_time_json='source_trading_time.json',
        target_trading_time_json='trading_time.json',
        offset_close_minutes=5,
    )