예제 소스

surge_redis_common.py

# -*- coding: utf-8 -*-
"""1분봉 OHLCV 를 Redis 에 JSON 으로 저장/로드 (writer ↔ consumer 공통).

키: ``{prefix}{종목코드}``  (기본 prefix ``hankook:surge:minute:``)

환경 변수 ``REDIS_URL`` (예: ``redis://127.0.0.1:6379/0``) 를 consumer/writer 가 공유합니다.
"""

from __future__ import annotations

import json
from datetime import datetime
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import pandas as pd

DEFAULT_MINUTE_KEY_PREFIX = "hankook:surge:minute:"


def minute_key(code: str, *, prefix: str = DEFAULT_MINUTE_KEY_PREFIX) -> str:
    return f"{prefix}{code}"


def bars_json_from_df(df: "pd.DataFrame") -> str:
    import pandas as pd

    bars: list[dict[str, Any]] = []
    for idx, row in df.iterrows():
        ts = idx if isinstance(idx, pd.Timestamp) else pd.Timestamp(idx)
        bars.append(
            {
                "t": ts.isoformat(),
                "O": float(row["Open"]),
                "H": float(row["High"]),
                "L": float(row["Low"]),
                "C": float(row["Close"]),
                "V": float(row["Volume"]),
            }
        )
    payload = {"updated_at": datetime.now().isoformat(timespec="seconds"), "bars": bars}
    return json.dumps(payload, ensure_ascii=False)


def df_from_redis_value(raw: str | bytes | None) -> tuple["pd.DataFrame", datetime | None]:
    """Redis GET 값 → (DataFrame index=DateTime, columns OHLCV), updated_at."""
    import pandas as pd

    if raw is None:
        return pd.DataFrame(columns=["Open", "High", "Low", "Close", "Volume"]), None
    if isinstance(raw, bytes):
        raw = raw.decode("utf-8")
    try:
        data = json.loads(raw)
    except Exception:
        return pd.DataFrame(columns=["Open", "High", "Low", "Close", "Volume"]), None
    upd_raw = data.get("updated_at")
    updated: datetime | None = None
    if upd_raw:
        try:
            updated = datetime.fromisoformat(str(upd_raw))
        except Exception:
            updated = None
    rows = data.get("bars") or []
    if not rows:
        return pd.DataFrame(columns=["Open", "High", "Low", "Close", "Volume"]), updated
    rec: list[dict[str, Any]] = []
    for b in rows:
        rec.append(
            {
                "DateTime": pd.Timestamp(b["t"]),
                "Open": float(b["O"]),
                "High": float(b["H"]),
                "Low": float(b["L"]),
                "Close": float(b["C"]),
                "Volume": float(b["V"]),
            }
        )
    out = pd.DataFrame(rec).set_index("DateTime").sort_index()
    return out[["Open", "High", "Low", "Close", "Volume"]], updated