okx-codex-trader/okx_codex_trader/time_rules.py

from __future__ import annotations

from zoneinfo import ZoneInfo

import pandas as pd


NY = ZoneInfo("America/New_York")
BJ = ZoneInfo("Asia/Shanghai")


def utc_timestamp(ts: int) -> pd.Timestamp:
    return pd.to_datetime(ts, unit="ms", utc=True)


def ny_minutes(ts: int) -> tuple[int, int]:
    dt = utc_timestamp(ts).to_pydatetime().astimezone(NY)
    return dt.weekday(), dt.hour * 60 + dt.minute


def is_ny_weekend(ts: int) -> bool:
    weekday, _ = ny_minutes(ts)
    return weekday >= 5


def is_us_open_window(ts: int) -> bool:
    weekday, minute = ny_minutes(ts)
    return weekday < 5 and 9 * 60 + 30 <= minute < 10 * 60 + 30


def is_us_open_extended(ts: int) -> bool:
    weekday, minute = ny_minutes(ts)
    return weekday < 5 and 9 * 60 <= minute < 11 * 60


def entry_allowed(ts: int, mode: str) -> bool:
    weekday, minute = ny_minutes(ts)
    if mode == "all":
        return True
    if mode == "weekday":
        return not is_ny_weekend(ts)
    if mode == "no_weekend":
        return not is_ny_weekend(ts)
    if mode == "weekend":
        return is_ny_weekend(ts)
    if mode == "no_us_open":
        return not is_us_open_extended(ts)
    if mode == "us_open_only":
        return is_us_open_extended(ts)
    if mode == "asia_bj":
        hour = utc_timestamp(ts).to_pydatetime().astimezone(BJ).hour
        return 8 <= hour < 16
    if mode == "us_regular":
        return weekday < 5 and 9 * 60 + 30 <= minute < 16 * 60
    raise ValueError("entry_time_filter is invalid")


def time_bucket(ts: int) -> str:
    if is_ny_weekend(ts):
        return "weekend"
    _, minute = ny_minutes(ts)
    if 4 * 60 <= minute < 9 * 60:
        return "us_premarket"
    if 9 * 60 <= minute < 9 * 60 + 30:
        return "us_preopen"
    if 9 * 60 + 30 <= minute < 10 * 60 + 30:
        return "us_open_1h"
    if 10 * 60 + 30 <= minute < 16 * 60:
        return "us_regular_late"
    if 16 * 60 <= minute < 20 * 60:
        return "us_afterhours"
    return "overnight"