File size: 3,708 Bytes
210f9af | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """KST buckets and session tags (pure Python).
Session tags (sentiment-labelling §6.2 hour bands; rollup column ``session``, classifier-plan §7.3):
``holiday`` on weekends and KRX holidays, else ``pre_open`` 06:00-09:00, ``regular`` 09:00-15:30,
``after_hours`` 15:30-20:00, ``night`` 20:00-06:00 (KST). An hourly bucket takes the tag of its start time, so
the 15:00 bucket (which holds the 15:30 close) is ``regular``. Holidays come from ``config/krx_holidays.yaml``
(a list of dates, or ``{"holidays": [...]}``); without it only weekends count.
"""
from __future__ import annotations
from datetime import date, datetime, time, timedelta, timezone
from pathlib import Path
from typing import Iterable
KST = timezone(timedelta(hours=9))
SESSIONS = ("pre_open", "regular", "after_hours", "night", "holiday")
def to_kst(dt: datetime) -> datetime:
if dt.tzinfo is None:
raise ValueError("naive datetime: written_at must carry a zone")
return dt.astimezone(KST)
def parse(text: str) -> datetime:
"""ISO time with an offset (``Z`` accepted) -> aware datetime."""
dt = datetime.fromisoformat(text.replace("Z", "+00:00"))
if dt.tzinfo is None:
raise ValueError(f"time without offset: {text!r}")
return dt
def hour_bucket(dt: datetime) -> datetime:
"""KST hour start as a naive local timestamp (``bucket_kst timestamp``)."""
k = to_kst(dt)
return k.replace(minute=0, second=0, microsecond=0, tzinfo=None)
def day_bucket(dt: datetime) -> date:
return to_kst(dt).date()
def is_trading_day(d: date, holidays: Iterable[date] = ()) -> bool:
return d.weekday() < 5 and d not in set(holidays)
def session_tag(dt: datetime, holidays: Iterable[date] = ()) -> str:
"""Session of an aware time, or of a naive KST timestamp (a bucket start)."""
k = dt.replace(tzinfo=None) if dt.tzinfo is None else to_kst(dt).replace(tzinfo=None)
if not is_trading_day(k.date(), holidays):
return "holiday"
t = k.time()
if time(6) <= t < time(9):
return "pre_open"
if time(9) <= t < time(15, 30):
return "regular"
if time(15, 30) <= t < time(20):
return "after_hours"
return "night"
def load_holidays(path: str | Path | None) -> set[date]:
if path is None or not Path(path).exists():
return set()
import yaml
value = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or []
if isinstance(value, dict):
value = value.get("holidays", [])
out = set()
for v in value:
if isinstance(v, dict):
v = v.get("date")
out.add(v if isinstance(v, date) else date.fromisoformat(str(v)))
return out
def iso_week(dt: datetime) -> str:
"""ISO week of the KST date, e.g. ``2026-W42``."""
y, w, _ = to_kst(dt).isocalendar()
return f"{y}-W{w:02d}"
def load_groups(path: str | Path | None, codes: Iterable[str] = ()) -> dict[str, str]:
"""Code -> stock group/sector from ``config/stocks.yaml`` (classifier-plan §5.2).
Accepted shapes: ``{"stocks": [{"code", "group"|"sector"}]}``, a list of such entries, or ``{code: group}``.
Codes without an entry (or without the file) are their own group."""
out = {c: c for c in codes}
if path is None or not Path(path).exists():
return out
import yaml
value = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}
entries = value.get("stocks", value) if isinstance(value, dict) else value
if isinstance(entries, dict):
out.update({str(k): str(v) for k, v in entries.items()})
return out
for e in entries:
code = str(e["code"])
out[code] = str(e.get("group") or e.get("sector") or code)
return out
|