1"""XNYS trading calendar helpers (America/New_York sessions, UTC storage)."""2from __future__ import annotations34import datetime as dt5from functools import lru_cache6from zoneinfo import ZoneInfo78import numpy as np9import pandas as pd1011ET = ZoneInfo("America/New_York")12UTC = ZoneInfo("UTC")131415@lru_cache(maxsize=1)16def _xnys():17 import exchange_calendars as xc18 return xc.get_calendar("XNYS", start="2000-01-01", end="2027-12-31")192021@lru_cache(maxsize=1)22def sessions() -> pd.DatetimeIndex:23 """All XNYS sessions 2000→2027 as tz-naive dates."""24 return _xnys().sessions.tz_localize(None) if _xnys().sessions.tz is not None else _xnys().sessions252627def session_array() -> np.ndarray:28 return sessions().to_numpy().astype("datetime64[D]")293031def is_session(day: dt.date) -> bool:32 return np.datetime64(day) in set(session_array().tolist())333435def shift_trading_days(day: dt.date | np.datetime64 | pd.Timestamp, n: int) -> dt.date:36 """Shift by n trading days; if `day` is not a session, start from the previous session (n<=0) or next (n>0)."""37 arr = session_array()38 d = np.datetime64(pd.Timestamp(day).date())39 pos = int(np.searchsorted(arr, d))40 if pos < len(arr) and arr[pos] == d:41 base = pos42 else:43 base = pos - 1 if n <= 0 else pos44 idx = base + n45 if idx < 0 or idx >= len(arr):46 raise ValueError("date out of calendar range")47 return arr[idx].astype("datetime64[D]").astype(dt.date)484950def trading_days_between(a: dt.date, b: dt.date) -> int:51 """Number of sessions strictly after a up to and including b (negative if b<a)."""52 arr = session_array()53 ia, ib = np.searchsorted(arr, np.datetime64(a), side="right"), np.searchsorted(arr, np.datetime64(b), side="right")54 return int(ib - ia)555657def month_end_sessions(start: dt.date, end: dt.date) -> list[dt.date]:58 s = pd.Series(sessions())59 s = s[(s >= pd.Timestamp(start)) & (s <= pd.Timestamp(end))]60 return [d.date() for d in s.groupby(s.dt.to_period("M")).max()]616263def close_time_utc(session: dt.date) -> pd.Timestamp:64 """Open instant (UTC) of the closing-print bar (16:00 ET, or 13:00 ET on early-close days)."""65 cal = _xnys()66 ts = pd.Timestamp(session)67 close = cal.session_close(ts)68 return pd.Timestamp(close).tz_convert("UTC") if close.tz is not None else pd.Timestamp(close).tz_localize("UTC")697071def regular_session_bounds_utc(session: dt.date) -> tuple[pd.Timestamp, pd.Timestamp]:72 """[open 09:30 ET, close bar open] in UTC for the session (inclusive bar-open bounds)."""73 cal = _xnys()74 ts = pd.Timestamp(session)75 o, c = cal.session_open(ts), cal.session_close(ts)76 o = pd.Timestamp(o).tz_convert("UTC") if o.tz is not None else pd.Timestamp(o).tz_localize("UTC")77 c = pd.Timestamp(c).tz_convert("UTC") if c.tz is not None else pd.Timestamp(c).tz_localize("UTC")78 return o, c79