stock-oracle/app/services/news/session_window.py

"""
Trading-session window math (NYSE / XNYS).

`session_window(date, "premarket"|"intraday"|"post"|"full_session")` returns
the (start_utc, end_utc) bounds the news aggregator filters on. NYSE holidays
and short-day closes (1pm ET on day-after-Thanksgiving etc.) are handled by
pandas_market_calendars when available; falls back to weekday-only logic if
the dependency is missing.
"""

from __future__ import annotations

import logging
from datetime import date, datetime, time, timedelta, timezone
from functools import lru_cache
from typing import Literal
from zoneinfo import ZoneInfo

logger = logging.getLogger(__name__)

ET = ZoneInfo("America/New_York")

WindowName = Literal["premarket", "intraday", "post", "full_session"]

# Default ET clock-time bounds. Short days (early close) override regular_close
# via the market-calendar lookup below. `post` window ends at the *next*
# trading day's premarket start (04:00 ET) to avoid overlap.
_DEFAULT_PREMARKET_START_T = time(4, 0)
_DEFAULT_REGULAR_OPEN_T = time(9, 30)
_DEFAULT_REGULAR_CLOSE_T = time(16, 0)


def session_window(session_date: date, window: WindowName) -> tuple[datetime, datetime]:
    """Return UTC bounds for the given window on the given ET session date."""
    if not _is_trading_day(session_date):
        raise ValueError(
            f"{session_date} is not a NYSE trading day; choose the next/prev session"
        )

    open_dt_et, close_dt_et = _session_bounds_et(session_date)
    prev_close_dt_et = _prev_session_close_et(session_date)
    next_premarket_start_et = _next_session_premarket_start_et(session_date)

    if window == "premarket":
        start_et = prev_close_dt_et
        end_et = open_dt_et
    elif window == "intraday":
        start_et = open_dt_et
        end_et = close_dt_et
    elif window == "post":
        # Post-market ends at next session's premarket start (04:00 ET) to avoid
        # overlap with the next session's `premarket` window.
        start_et = close_dt_et
        end_et = next_premarket_start_et
    elif window == "full_session":
        start_et = prev_close_dt_et
        end_et = next_premarket_start_et
    else:
        raise ValueError(f"Unknown window: {window}")

    return _to_utc(start_et), _to_utc(end_et)


def _to_utc(dt: datetime) -> datetime:
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=ET)
    return dt.astimezone(timezone.utc)


# ---------------------------------------------------------------------------
# Market-calendar lookups (cached)
# ---------------------------------------------------------------------------

@lru_cache(maxsize=1)
def _calendar():
    try:
        import pandas_market_calendars as mcal  # type: ignore
        return mcal.get_calendar("XNYS")
    except Exception as e:
        logger.warning(
            "pandas_market_calendars not available (%s); falling back to weekday-only session logic",
            e,
        )
        return None


def _is_trading_day(d: date) -> bool:
    cal = _calendar()
    if cal is None:
        return d.weekday() < 5
    schedule = cal.schedule(start_date=d, end_date=d)
    return not schedule.empty


def _session_bounds_et(d: date) -> tuple[datetime, datetime]:
    """Regular open/close in ET. Honors short days when calendar is available."""
    cal = _calendar()
    if cal is not None:
        schedule = cal.schedule(start_date=d, end_date=d)
        if not schedule.empty:
            row = schedule.iloc[0]
            open_utc = row["market_open"].to_pydatetime()
            close_utc = row["market_close"].to_pydatetime()
            return open_utc.astimezone(ET), close_utc.astimezone(ET)
    return (
        datetime.combine(d, _DEFAULT_REGULAR_OPEN_T, tzinfo=ET),
        datetime.combine(d, _DEFAULT_REGULAR_CLOSE_T, tzinfo=ET),
    )


def _prev_session_close_et(d: date) -> datetime:
    prev = _prev_trading_day(d)
    _, close_dt_et = _session_bounds_et(prev)
    return close_dt_et


def _next_session_open_et(d: date) -> datetime:
    nxt = _next_trading_day(d)
    open_dt_et, _ = _session_bounds_et(nxt)
    return open_dt_et


def _next_session_premarket_start_et(d: date) -> datetime:
    """Next trading day's premarket window start (04:00 ET)."""
    nxt = _next_trading_day(d)
    return datetime.combine(nxt, _DEFAULT_PREMARKET_START_T, tzinfo=ET)


def _prev_trading_day(d: date, max_lookback: int = 10) -> date:
    cal = _calendar()
    if cal is not None:
        schedule = cal.schedule(start_date=d - timedelta(days=max_lookback), end_date=d - timedelta(days=1))
        if not schedule.empty:
            return schedule.index[-1].date()
    candidate = d - timedelta(days=1)
    while candidate.weekday() >= 5:
        candidate -= timedelta(days=1)
    return candidate


def _next_trading_day(d: date, max_lookahead: int = 10) -> date:
    cal = _calendar()
    if cal is not None:
        schedule = cal.schedule(start_date=d + timedelta(days=1), end_date=d + timedelta(days=max_lookahead))
        if not schedule.empty:
            return schedule.index[0].date()
    candidate = d + timedelta(days=1)
    while candidate.weekday() >= 5:
        candidate += timedelta(days=1)
    return candidate