브랜치 분리 방식: A / B / C

A 선택 시 커밋 메시지: 위 초안 OK / 수정 / 직접 작성
작업 시점: 지금 / 운영 데이터 1~2일 쌓고 / 주말

kiwwom_trader/ws_candle_agg.py

@@ -0,0 +1,122 @@
+# 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