kis_trader/remove/legacy_root/kiwwom_trader/ws_candle_agg.py

# kiwwom_trader/ws_candle_agg.py
# 키움 웹소켓 실시간 체결(0B) 수신 → 1분봉 집계 → DB 또는 메모리 버퍼

import asyncio
import logging
from collections import defaultdict
from datetime import datetime
from typing import Callable, Dict, List, Optional

logger = logging.getLogger(__name__)


def _slot_1m(now: datetime) -> str:
    """현재 시각 기준 1분 봉 키 (YYYYMMDDHHMI)."""
    return now.strftime("%Y%m%d%H%M")


class KiwoomCandleAggregator:
    """
    키움 실시간 타입 0B(주식체결) 수신 시 호출되는 콜백에서
    종목별·1분봉별 OHLCV를 집계. 확정된 봉은 on_candle 콜백 또는 DB에 전달.
    """
    def __init__(
        self,
        on_candle: Optional[Callable[[str, int, str, dict], None]] = None,
        timeframe_min: int = 1,
    ):
        """
        Args:
            on_candle: (code, timeframe_min, candle_time, candle_dict) 호출. DB 저장 시 사용.
            timeframe_min: 봉 간격(분). 1 = 1분봉.
        """
        self.on_candle = on_candle
        self.timeframe_min = timeframe_min
        # code -> slot_1m -> { open, high, low, close, volume }
        self._buf: Dict[str, Dict[str, dict]] = defaultdict(dict)
        self._lock = asyncio.Lock()

    def push_tick(self, code: str, price: float, volume: int, tick_time: Optional[datetime] = None):
        """
        체결 1건 누적. 스레드 세이프하지 않음; 이벤트 루프 한 스레드에서 호출 권장.
        """
        t = tick_time or datetime.now()
        slot = _slot_1m(t)
        if slot not in self._buf[code]:
            self._buf[code][slot] = {
                "open": price,
                "high": price,
                "low": price,
                "close": price,
                "volume": volume,
            }
        else:
            c = self._buf[code][slot]
            c["high"] = max(c["high"], price)
            c["low"] = min(c["low"], price)
            c["close"] = price
            c["volume"] = c.get("volume", 0) + volume

    def flush_slot(self, code: str, slot: str) -> Optional[dict]:
        """해당 종목·슬롯 봉을 확정해 반환하고 버퍼에서 제거."""
        if code not in self._buf or slot not in self._buf[code]:
            return None
        c = self._buf[code].pop(slot)
        if not self._buf[code]:
            del self._buf[code]
        return c

    def get_candles(self, code: str, timeframe_min: int, n: int = 50) -> List[dict]:
        """
        메모리 버퍼만 사용하는 경우: 현재 버퍼에 쌓인 봉 반환. 없으면 빈 리스트.
        DB와 연동 시에는 이 메서드에서 DB를 조회하도록 확장.
        """
        if code not in self._buf:
            return []
        slots = sorted(self._buf[code].keys())
        out = []
        for slot in slots[-n:]:
            c = self._buf[code][slot].copy()
            c["candle_time"] = slot
            c["volume"] = c.get("volume", 0)
            out.append(c)
        return out


def build_ws_client_with_candle_callback(
    access_token: str,
    aggregator: KiwoomCandleAggregator,
    ws_url: Optional[str] = None,
):
    """
    키움 WebSocketClient에 실시간 0B 수신 시 aggregator.push_tick 호출하도록 연결.
    kiwoom_rest_api.websocket / websocket_helper 사용.
    """
    try:
        from kiwoom_rest_api.websocket import WebSocketClient, RealTimeData
        from kiwoom_rest_api.websocket_constants import STOCK_TRADE_FIELDS, get_field_name
    except ImportError:
        logger.warning("kiwoom_rest_api 미설치 또는 경로 미등록 → 웹소켓 캔들 비활성")
        return None

    client = WebSocketClient(access_token=access_token, ws_url=ws_url)

    async def on_data(realtime_data: RealTimeData):
        if realtime_data.trnm != "REAL":
            return
        for item in realtime_data.data or []:
            type_code = item.get("type", "")
            if type_code != "0B":
                continue
            item_code = (item.get("item") or [""])[0] if isinstance(item.get("item"), list) else item.get("item", "")
            if not item_code:
                continue
            values = item.get("values") or {}
            # 10:현재가, 15:거래량 등 (STOCK_TRADE_FIELDS)
            price = float(values.get("10", 0) or 0)
            vol = int(values.get("15", 0) or 0)
            if price > 0:
                aggregator.push_tick(item_code, price, vol)

    client.on_data = on_data
    return client