182 lines
7.7 KiB
Python
182 lines
7.7 KiB
Python
"""Exchange-calendar wrapper: US/KRX sessions, close times and briefing schedule.
|
|
|
|
All date arithmetic goes through `exchange_calendars` (XNYS / XKRX) and `zoneinfo`; no UTC
|
|
offsets are hardcoded, so DST switches, holidays and early closes are handled automatically.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections.abc import Iterable
|
|
from datetime import date, datetime, timedelta
|
|
from zoneinfo import ZoneInfo
|
|
|
|
import exchange_calendars as xc
|
|
import pandas as pd
|
|
|
|
from us2kr.config import ET, KST
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
ET_TZ = ZoneInfo(ET)
|
|
KST_TZ = ZoneInfo(KST)
|
|
|
|
_US_CAL = "XNYS"
|
|
_KR_CAL = "XKRX"
|
|
|
|
# KRX closures announced by KRX / the Korean government but missing from the pinned exchange_calendars
|
|
# data (4.13.2). Election days and 임시공휴일 are decided after a library release, so the data lags.
|
|
# Drop an entry once a release marks it as a holiday (the override is harmless while both agree).
|
|
KNOWN_KRX_CLOSURES: frozenset[date] = frozenset(
|
|
{
|
|
date(2025, 10, 10), # 임시공휴일 (추석 연휴 연장)
|
|
date(2026, 6, 3), # 제9회 전국동시지방선거
|
|
}
|
|
)
|
|
|
|
|
|
class MarketClock:
|
|
"""Session / briefing-time calculator backed by XNYS and XKRX calendars.
|
|
|
|
Calendars are built once for a window around `anchor` (default: today) and rebuilt
|
|
transparently if a query falls outside that window (long-running daemon safety).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
run_offset_minutes: int = 75,
|
|
*,
|
|
anchor: date | None = None,
|
|
years_back: int = 2,
|
|
years_ahead: int = 2,
|
|
extra_krx_holidays: Iterable[date] | None = None,
|
|
) -> None:
|
|
if run_offset_minutes < 0:
|
|
raise ValueError("run_offset_minutes must be >= 0")
|
|
self.run_offset = timedelta(minutes=run_offset_minutes)
|
|
self._years_back = years_back
|
|
self._years_ahead = years_ahead
|
|
# Ad-hoc KRX closures on top of the library data (KNOWN_KRX_CLOSURES + operator-supplied dates).
|
|
self._krx_closures: frozenset[date] = KNOWN_KRX_CLOSURES | frozenset(extra_krx_holidays or ())
|
|
self._build(anchor or datetime.now(tz=KST_TZ).date())
|
|
|
|
# ------------------------------------------------------------------ calendars
|
|
|
|
# Pad the requested range so Jan 1 .. Dec 31 of the window always lie strictly inside
|
|
# [first_session, last_session] of both calendars (no exchange closes > ~5 days in a row).
|
|
_PAD = timedelta(days=10)
|
|
|
|
def _build(self, anchor: date) -> None:
|
|
start = date(anchor.year - self._years_back, 1, 1) - self._PAD
|
|
end = date(anchor.year + self._years_ahead, 12, 31) + self._PAD
|
|
self._xnys = xc.get_calendar(_US_CAL, start=start, end=end)
|
|
self._xkrx = xc.get_calendar(_KR_CAL, start=start, end=end)
|
|
# Bounds common to both calendars (sessions are tz-naive Timestamps in v4).
|
|
self._first = max(self._xnys.first_session, self._xkrx.first_session).date()
|
|
self._last = min(self._xnys.last_session, self._xkrx.last_session).date()
|
|
log.debug("calendars built: %s .. %s", self._first, self._last)
|
|
|
|
@staticmethod
|
|
def _as_date(d: date | datetime) -> date:
|
|
return d.date() if isinstance(d, datetime) else d
|
|
|
|
def _ts(self, d: date | datetime) -> pd.Timestamp:
|
|
"""Normalise to a tz-naive midnight Timestamp inside the calendar range.
|
|
|
|
May rebuild `self._xnys` / `self._xkrx`, so callers must evaluate this BEFORE
|
|
binding a calendar method (`ts = self._ts(d); self._xnys.f(ts)`).
|
|
"""
|
|
d = self._as_date(d)
|
|
if not (self._first <= d <= self._last):
|
|
log.info("date %s outside calendar window %s..%s; rebuilding", d, self._first, self._last)
|
|
self._build(d)
|
|
return pd.Timestamp(d)
|
|
|
|
@staticmethod
|
|
def _require_aware(now: datetime) -> datetime:
|
|
if now.tzinfo is None or now.utcoffset() is None:
|
|
raise ValueError("`now` must be a tz-aware datetime")
|
|
return now
|
|
|
|
# ------------------------------------------------------------------ sessions
|
|
|
|
def is_us_session(self, d: date) -> bool:
|
|
"""True if XNYS trades on calendar date `d`."""
|
|
ts = self._ts(d)
|
|
return bool(self._xnys.is_session(ts))
|
|
|
|
def is_krx_session(self, d: date) -> bool:
|
|
"""True if KRX trades on calendar date `d` (library data + known ad-hoc closures)."""
|
|
if self._as_date(d) in self._krx_closures:
|
|
return False
|
|
ts = self._ts(d)
|
|
return bool(self._xkrx.is_session(ts))
|
|
|
|
def next_krx_session(self, on_or_after: date) -> date:
|
|
"""Earliest KRX session on or after `on_or_after` (ad-hoc closures respected)."""
|
|
d = self._as_date(on_or_after)
|
|
while True:
|
|
ts = self._ts(d)
|
|
d = self._xkrx.date_to_session(ts, direction="next").date()
|
|
if d not in self._krx_closures:
|
|
return d
|
|
d += timedelta(days=1)
|
|
|
|
def us_session_close(self, session: date) -> datetime:
|
|
"""Regular-session close as tz-aware America/New_York datetime (early closes respected)."""
|
|
ts = self._ts(session)
|
|
if not self._xnys.is_session(ts):
|
|
raise ValueError(f"{self._as_date(session)} is not an XNYS session")
|
|
close_utc: pd.Timestamp = self._xnys.session_close(ts)
|
|
return close_utc.to_pydatetime().astimezone(ET_TZ)
|
|
|
|
def previous_us_session(self, before: date) -> date:
|
|
"""Latest US session strictly before `before`."""
|
|
ts = self._ts(self._as_date(before) - timedelta(days=1))
|
|
return self._xnys.date_to_session(ts, direction="previous").date()
|
|
|
|
def next_us_session(self, after: date) -> date:
|
|
"""Earliest US session strictly after `after`."""
|
|
ts = self._ts(self._as_date(after) + timedelta(days=1))
|
|
return self._xnys.date_to_session(ts, direction="next").date()
|
|
|
|
def latest_completed_session(self, now: datetime) -> date:
|
|
"""Latest US session whose regular close is <= `now`."""
|
|
now = self._require_aware(now)
|
|
today_et = now.astimezone(ET_TZ).date()
|
|
ts = self._ts(today_et)
|
|
candidate = self._xnys.date_to_session(ts, direction="previous").date()
|
|
while self.us_session_close(candidate) > now:
|
|
candidate = self.previous_us_session(candidate)
|
|
return candidate
|
|
|
|
# ------------------------------------------------------------------ briefing schedule
|
|
|
|
def briefing_time_for(self, session: date) -> datetime:
|
|
"""Briefing run time = session close + run offset, expressed in Asia/Seoul."""
|
|
return (self.us_session_close(session) + self.run_offset).astimezone(KST_TZ)
|
|
|
|
def krx_date_for(self, session: date) -> date:
|
|
"""Korean calendar date on which the briefing for `session` is delivered."""
|
|
return self.briefing_time_for(session).date()
|
|
|
|
def next_briefing(self, now: datetime) -> tuple[date, datetime]:
|
|
"""(session, run_at KST) for the earliest briefing strictly after `now`."""
|
|
now = self._require_aware(now)
|
|
today_et = now.astimezone(ET_TZ).date()
|
|
session = self.previous_us_session(today_et)
|
|
while self.briefing_time_for(session) <= now:
|
|
session = self.next_us_session(session)
|
|
return session, self.briefing_time_for(session)
|
|
|
|
def catchup_candidate(self, now: datetime, window_hours: int) -> date | None:
|
|
"""Session whose briefing time lies in (now - window, now]; None when nothing was missed."""
|
|
now = self._require_aware(now)
|
|
window = timedelta(hours=window_hours)
|
|
session = self.latest_completed_session(now)
|
|
while True:
|
|
run_at = self.briefing_time_for(session)
|
|
if run_at <= now:
|
|
return session if run_at > now - window else None
|
|
session = self.previous_us_session(session)
|