feat(tests): 신규 키움 웹소켓 조건검색 및 실시간 조건검색 테스트 추가
변경 사항 ---- - _test_kiwoom_condition_list.py: 키움 웹소켓 조건검색 '목록조회' 기능을 단독으로 테스트하는 스크립트 추가 - _test_kiwoom_condition_realtime.py: 'momentum' 조건식을 실시간으로 등록하고 초기 매칭 종목 리스트 및 실시간 편입/이탈을 수신하는 테스트 스크립트 추가 - _verify_columnar_bitid.py, _verify_shared_e2e_breakout.py, _verify_shared_e2e.py: 공유 메모리 및 dict 간의 데이터 일관성을 검증하는 테스트 추가 영향 ---- - 신규 테스트 스크립트 추가로 키움 웹소켓 API의 기능 검증 및 안정성을 높임 - 기존 기능에 대한 영향 없음 Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
412
kis_trader/engine/indicator_cache.py
Normal file
412
kis_trader/engine/indicator_cache.py
Normal file
@@ -0,0 +1,412 @@
|
||||
"""
|
||||
kis_trader/engine/indicator_cache.py — RSI/EMA materialized 캐시 (파라서치 가속)
|
||||
================================================================================
|
||||
- 종목별 IndicatorCache: RSI·EMA period 별 시리즈 1회 계산 → 조합 루프 재사용
|
||||
- ws_candles DB 컬럼 materialize (RSI·EMA period 별 rsi_N / ema_N)
|
||||
|
||||
[ws_candles 컬럼]
|
||||
- 실시간 WS 확정봉: rsi_2, rsi_3, rsi_5 (CandleAggregator 가 INSERT)
|
||||
- materialize 배치: rsi_7, rsi_14, rsi_21 + ema_5,9,12,15,21,34 (env 로 period 확장)
|
||||
|
||||
파라서치 ④ 구간(조합×RSI/EMA 재계산) 가속용 — 실매 TRIGGER 로직 변경 없음.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
||||
|
||||
from kis_trader.engine.ema_trend_filter import compute_ema_series
|
||||
from kis_trader.utils.env import get_env_bool, get_env_float, get_env_from_db, get_env_int
|
||||
|
||||
|
||||
def _parse_period_list(env_key: str, fallback: str) -> List[int]:
|
||||
raw = str(get_env_from_db(env_key, fallback) or fallback).strip()
|
||||
out: List[int] = []
|
||||
for part in raw.split(","):
|
||||
part = part.strip()
|
||||
if not part:
|
||||
continue
|
||||
try:
|
||||
p = int(float(part))
|
||||
if p > 0:
|
||||
out.append(p)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
return sorted(set(out)) if out else [int(x) for x in fallback.split(",") if x.strip()]
|
||||
|
||||
|
||||
def default_rsi_periods() -> List[int]:
|
||||
return _parse_period_list("PARAM_SEARCH_RSI_PERIODS", "2,3,5,7,14,21")
|
||||
|
||||
|
||||
def default_ema_periods() -> List[int]:
|
||||
return _parse_period_list("PARAM_SEARCH_EMA_PERIODS", "5,9,12,15,21,34")
|
||||
|
||||
|
||||
def ws_candles_materialize_rsi_periods() -> List[int]:
|
||||
"""DB materialize 대상 RSI period — env ``WS_CANDLE_MATERIALIZE_RSI_PERIODS`` 우선."""
|
||||
raw = _parse_period_list(
|
||||
"WS_CANDLE_MATERIALIZE_RSI_PERIODS",
|
||||
",".join(str(p) for p in default_rsi_periods() if p not in (2, 3, 5)),
|
||||
)
|
||||
# rsi_2/3/5 는 CandleAggregator 가 실시간 기록 — 배치는 7+ 위주
|
||||
return sorted(set(raw))
|
||||
|
||||
|
||||
def ws_candles_materialize_ema_periods() -> List[int]:
|
||||
"""DB materialize 대상 EMA period."""
|
||||
return _parse_period_list(
|
||||
"WS_CANDLE_MATERIALIZE_EMA_PERIODS",
|
||||
",".join(str(p) for p in default_ema_periods()),
|
||||
)
|
||||
|
||||
|
||||
def ws_candles_all_rsi_column_names() -> List[str]:
|
||||
"""ws_candles SELECT/마이그레이션용 RSI 컬럼명 (rsi_2 … env period)."""
|
||||
periods = sorted(set(default_rsi_periods()) | {2, 3, 5})
|
||||
return [f"rsi_{p}" for p in periods]
|
||||
|
||||
|
||||
def ws_candles_all_ema_column_names() -> List[str]:
|
||||
periods = ws_candles_materialize_ema_periods()
|
||||
return [f"ema_{p}" for p in periods]
|
||||
|
||||
|
||||
def ws_candles_indicator_column_migrations() -> List[Tuple[str, str]]:
|
||||
"""(컬럼명, ALTER DDL suffix) — TradeDB._migrate_add_columns 에서 사용.
|
||||
|
||||
※ TradeDB 기동 중 호출되므로 get_env_from_db(DB) 를 쓰지 않는다.
|
||||
(env 조회 → TradeDB() 재진입 → _tables_lock 데드락 → journal 로그 0줄 hang)
|
||||
"""
|
||||
out: List[Tuple[str, str]] = []
|
||||
# env 기본값과 동일한 정적 period — 런타임 확장 period 는 ensure 시 information_schema 로 ADD
|
||||
migration_rsi = (7, 14, 21)
|
||||
migration_ema = (5, 9, 12, 15, 21, 34)
|
||||
for period in migration_rsi:
|
||||
out.append((
|
||||
f"rsi_{period}",
|
||||
f"DOUBLE NULL COMMENT 'RSI({period}) materialized'",
|
||||
))
|
||||
for period in migration_ema:
|
||||
out.append((
|
||||
f"ema_{period}",
|
||||
f"DOUBLE NULL COMMENT 'EMA({period}) materialized'",
|
||||
))
|
||||
return out
|
||||
|
||||
|
||||
def ensure_ws_candles_indicator_columns(db) -> int:
|
||||
"""ws_candles 에 누락된 rsi_N / ema_N 컬럼 ADD (1회 마이그레이션). Returns 추가된 컬럼 수."""
|
||||
added = 0
|
||||
try:
|
||||
cols = set(db.conn.get_columns("ws_candles"))
|
||||
except Exception:
|
||||
return 0
|
||||
for col, ddl in ws_candles_indicator_column_migrations():
|
||||
if col in cols:
|
||||
continue
|
||||
try:
|
||||
db.conn.execute(f"ALTER TABLE ws_candles ADD COLUMN {col} {ddl}")
|
||||
cols.add(col)
|
||||
added += 1
|
||||
except Exception:
|
||||
pass
|
||||
return added
|
||||
|
||||
|
||||
def indicator_cache_enabled() -> bool:
|
||||
return get_env_bool("PARAM_SEARCH_INDICATOR_CACHE_ENABLED", True)
|
||||
|
||||
|
||||
def materialize_db_on_load() -> bool:
|
||||
return get_env_bool("WS_CANDLE_MATERIALIZE_ON_LOAD", False)
|
||||
|
||||
|
||||
class IndicatorCache:
|
||||
"""단일 종목 캔들에 대한 RSI/EMA 시리즈 메모이제이션."""
|
||||
|
||||
# _rsi_final / _ema_final: 해당 period 시리즈를 (DB seed + 1회 계산 병합 후)
|
||||
# '확정'했음을 표시. 확정된 period 는 재계산 없이 캐시를 그대로 반환한다.
|
||||
# (기존엔 rsi_at/ema_at 호출마다 compute_*_series 를 다시 돌려 O(N²) 였음 — 결과는
|
||||
# 동일했으나 파라서치 단일 백테가 봉당 전체 재계산으로 수십 초 소요되던 병목)
|
||||
__slots__ = ("closes", "_rsi", "_ema", "_rsi_final", "_ema_final")
|
||||
|
||||
def __init__(self, candles: List[Dict[str, Any]]) -> None:
|
||||
self.closes: List[float] = [float(c.get("close") or 0) for c in candles]
|
||||
self._rsi: Dict[int, List[Optional[float]]] = {}
|
||||
self._ema: Dict[int, List[Optional[float]]] = {}
|
||||
self._rsi_final: set = set()
|
||||
self._ema_final: set = set()
|
||||
self._seed_from_candle_fields(candles)
|
||||
|
||||
def _seed_from_candle_fields(self, candles: List[Dict[str, Any]]) -> None:
|
||||
"""DB materialized 컬럼(rsi_14, ema_9 등)이 있으면 해당 period 시리즈에 반영."""
|
||||
for i, c in enumerate(candles):
|
||||
for period in default_rsi_periods():
|
||||
key = f"rsi_{period}"
|
||||
if key in c and c[key] is not None:
|
||||
ser = self._rsi.setdefault(period, [None] * len(candles))
|
||||
if i < len(ser):
|
||||
try:
|
||||
ser[i] = float(c[key])
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
legacy = {2: "rsi_2", 3: "rsi_3", 5: "rsi_5"}.get(period)
|
||||
if legacy and legacy in c and c[legacy] is not None:
|
||||
ser = self._rsi.setdefault(period, [None] * len(candles))
|
||||
if i < len(ser) and ser[i] is None:
|
||||
try:
|
||||
ser[i] = float(c[legacy])
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
for period in default_ema_periods():
|
||||
key = f"ema_{period}"
|
||||
if key in c and c[key] is not None:
|
||||
ser = self._ema.setdefault(period, [None] * len(candles))
|
||||
if i < len(ser):
|
||||
try:
|
||||
ser[i] = float(c[key])
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
|
||||
def rsi_series(self, period: int) -> List[Optional[float]]:
|
||||
p = max(1, int(period))
|
||||
# 이미 확정된 period 면 재계산 없이 캐시 반환 (핫패스 가속, 결과 불변)
|
||||
if p in self._rsi_final:
|
||||
return self._rsi[p]
|
||||
from kis_trader.engine.momentum_engine import compute_rsi_series
|
||||
computed = compute_rsi_series(self.closes, p)
|
||||
seeded = self._rsi.get(p)
|
||||
if seeded is None:
|
||||
self._rsi[p] = computed
|
||||
self._rsi_final.add(p)
|
||||
return computed
|
||||
if len(seeded) != len(computed):
|
||||
seeded = (list(seeded) + [None] * len(computed))[: len(computed)]
|
||||
# warm-up prepend 등으로 DB rsi_N 이 앞구간만 있으면 뒤 구간은 계산값으로 보완
|
||||
merged = [
|
||||
s if s is not None else b
|
||||
for s, b in zip(seeded, computed)
|
||||
]
|
||||
self._rsi[p] = merged
|
||||
self._rsi_final.add(p)
|
||||
return merged
|
||||
|
||||
def rsi_at(self, i: int, period: int) -> Optional[float]:
|
||||
ser = self.rsi_series(period)
|
||||
if 0 <= i < len(ser):
|
||||
return ser[i]
|
||||
return None
|
||||
|
||||
def ema_series(self, period: int) -> List[Optional[float]]:
|
||||
p = max(1, int(period))
|
||||
# 이미 확정된 period 면 재계산 없이 캐시 반환 (핫패스 가속, 결과 불변)
|
||||
if p in self._ema_final:
|
||||
return self._ema[p]
|
||||
computed = compute_ema_series(self.closes, p)
|
||||
seeded = self._ema.get(p)
|
||||
if seeded is None:
|
||||
self._ema[p] = computed
|
||||
self._ema_final.add(p)
|
||||
return computed
|
||||
if len(seeded) != len(computed):
|
||||
seeded = (list(seeded) + [None] * len(computed))[: len(computed)]
|
||||
merged = [
|
||||
s if s is not None else b
|
||||
for s, b in zip(seeded, computed)
|
||||
]
|
||||
self._ema[p] = merged
|
||||
self._ema_final.add(p)
|
||||
return merged
|
||||
|
||||
def ema_at(self, i: int, period: int) -> Optional[float]:
|
||||
ser = self.ema_series(period)
|
||||
if 0 <= i < len(ser):
|
||||
return ser[i]
|
||||
return None
|
||||
|
||||
|
||||
def build_indicator_cache(candles: List[Dict[str, Any]]) -> IndicatorCache:
|
||||
return IndicatorCache(candles)
|
||||
|
||||
|
||||
def build_indicator_cache_by_code(
|
||||
codes_candles: Dict[str, List[Dict[str, Any]]],
|
||||
*,
|
||||
rsi_periods: Optional[Sequence[int]] = None,
|
||||
ema_periods: Optional[Sequence[int]] = None,
|
||||
) -> Dict[str, IndicatorCache]:
|
||||
"""종목별 캐시 — rsi/ema period 는 warm-up 용 (첫 접근 시 lazy 계산)."""
|
||||
if not indicator_cache_enabled():
|
||||
return {}
|
||||
out: Dict[str, IndicatorCache] = {}
|
||||
rsi_p = list(rsi_periods or default_rsi_periods())
|
||||
ema_p = list(ema_periods or default_ema_periods())
|
||||
for code, rows in (codes_candles or {}).items():
|
||||
if not rows:
|
||||
continue
|
||||
ic = IndicatorCache([dict(r) for r in rows])
|
||||
for p in rsi_p:
|
||||
ic.rsi_series(p)
|
||||
for p in ema_p:
|
||||
ic.ema_series(p)
|
||||
out[str(code)] = ic
|
||||
return out
|
||||
|
||||
|
||||
def attach_indicator_caches_to_params(
|
||||
params: Dict[str, Any],
|
||||
codes_candles: Dict[str, List[Dict[str, Any]]],
|
||||
) -> None:
|
||||
"""params['_indicator_cache_by_code'] 에 종목별 캐시 주입 (in-place)."""
|
||||
if not indicator_cache_enabled():
|
||||
return
|
||||
existing = params.get("_indicator_cache_by_code")
|
||||
if isinstance(existing, dict) and existing:
|
||||
return
|
||||
params["_indicator_cache_by_code"] = build_indicator_cache_by_code(codes_candles)
|
||||
|
||||
|
||||
def get_indicator_cache_from_params(
|
||||
params: Dict[str, Any],
|
||||
code: str,
|
||||
) -> Optional[IndicatorCache]:
|
||||
by_code = params.get("_indicator_cache_by_code")
|
||||
if isinstance(by_code, dict):
|
||||
return by_code.get(str(code))
|
||||
ic = params.get("_indicator_cache")
|
||||
return ic if isinstance(ic, IndicatorCache) else None
|
||||
|
||||
|
||||
def enrich_candles_with_materialized_fields(
|
||||
candles: List[Dict[str, Any]],
|
||||
ic: IndicatorCache,
|
||||
*,
|
||||
rsi_periods: Optional[Sequence[int]] = None,
|
||||
ema_periods: Optional[Sequence[int]] = None,
|
||||
) -> None:
|
||||
"""캔들 dict 에 rsi_N / ema_N 필드 부착 (DB 저장·재로드용)."""
|
||||
rsi_p = list(rsi_periods or default_rsi_periods())
|
||||
ema_p = list(ema_periods or default_ema_periods())
|
||||
for i, c in enumerate(candles):
|
||||
for p in rsi_p:
|
||||
v = ic.rsi_at(i, p)
|
||||
if v is not None:
|
||||
c[f"rsi_{p}"] = round(v, 4)
|
||||
for p in ema_p:
|
||||
v = ic.ema_at(i, p)
|
||||
if v is not None:
|
||||
c[f"ema_{p}"] = round(v, 4)
|
||||
|
||||
|
||||
def materialize_ws_candles_batch(
|
||||
db,
|
||||
candles_by_code: Dict[str, List[Dict[str, Any]]],
|
||||
timeframe: int,
|
||||
) -> int:
|
||||
"""
|
||||
ws_candles 행에 rsi_N·ema_N UPDATE (배치).
|
||||
env ``WS_CANDLE_MATERIALIZE_ON_LOAD=true`` 일 때만 실행.
|
||||
Returns: 업데이트 시도 행 수.
|
||||
"""
|
||||
if not materialize_db_on_load() or db is None:
|
||||
return 0
|
||||
ensure_ws_candles_indicator_columns(db)
|
||||
try:
|
||||
db_cols = set(db.conn.get_columns("ws_candles"))
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
rsi_p = [p for p in ws_candles_materialize_rsi_periods() if f"rsi_{p}" in db_cols]
|
||||
ema_p = [p for p in ws_candles_materialize_ema_periods() if f"ema_{p}" in db_cols]
|
||||
if not rsi_p and not ema_p:
|
||||
return 0
|
||||
|
||||
set_parts = []
|
||||
for p in rsi_p:
|
||||
set_parts.append(f"rsi_{p} = ?")
|
||||
for p in ema_p:
|
||||
set_parts.append(f"ema_{p} = ?")
|
||||
set_sql = ", ".join(set_parts)
|
||||
if not set_sql:
|
||||
return 0
|
||||
|
||||
batch: List[Tuple[Any, ...]] = []
|
||||
for code, rows in (candles_by_code or {}).items():
|
||||
if not rows:
|
||||
continue
|
||||
ic = build_indicator_cache([dict(r) for r in rows])
|
||||
enrich_candles_with_materialized_fields(
|
||||
rows, ic,
|
||||
rsi_periods=rsi_p or None,
|
||||
ema_periods=ema_p or None,
|
||||
)
|
||||
for c in rows:
|
||||
ct = str(c.get("candle_time") or "")
|
||||
if not ct:
|
||||
continue
|
||||
vals: List[Any] = []
|
||||
for p in rsi_p:
|
||||
vals.append(c.get(f"rsi_{p}"))
|
||||
for p in ema_p:
|
||||
vals.append(c.get(f"ema_{p}"))
|
||||
vals.extend([str(code), int(timeframe), ct])
|
||||
batch.append(tuple(vals))
|
||||
|
||||
if not batch:
|
||||
return 0
|
||||
|
||||
where = "WHERE code = ? AND timeframe = ? AND candle_time = ?"
|
||||
sql_sqlite = f"""
|
||||
UPDATE ws_candles
|
||||
SET {set_sql}, updated_at = datetime('now','localtime')
|
||||
{where}
|
||||
"""
|
||||
set_mysql = ", ".join(
|
||||
[f"rsi_{p} = %s" for p in rsi_p]
|
||||
+ [f"ema_{p} = %s" for p in ema_p]
|
||||
)
|
||||
sql_mysql = f"""
|
||||
UPDATE ws_candles
|
||||
SET {set_mysql}
|
||||
WHERE code = %s AND timeframe = %s AND candle_time = %s
|
||||
"""
|
||||
try:
|
||||
with db.conn:
|
||||
db.conn.executemany(sql_sqlite, batch)
|
||||
except Exception:
|
||||
try:
|
||||
db.conn.executemany(sql_mysql, batch)
|
||||
except Exception:
|
||||
return 0
|
||||
return len(batch)
|
||||
|
||||
|
||||
def ws_candles_select_indicator_cols(db=None) -> str:
|
||||
"""SELECT 절 추가 컬럼 — DB 에 실제 존재하는 rsi_N / ema_N 만."""
|
||||
names: List[str] = []
|
||||
if db is not None:
|
||||
try:
|
||||
cols = set(db.conn.get_columns("ws_candles"))
|
||||
for n in ws_candles_all_rsi_column_names() + ws_candles_all_ema_column_names():
|
||||
if n in cols:
|
||||
names.append(n)
|
||||
except Exception:
|
||||
pass
|
||||
if not names:
|
||||
names = ws_candles_all_rsi_column_names() + ws_candles_all_ema_column_names()
|
||||
return (", " + ", ".join(names)) if names else ""
|
||||
|
||||
|
||||
def ws_candles_has_materialized_cols(db) -> bool:
|
||||
"""materialize 컬럼(ema_9 또는 ema_34 등)이 1개 이상 있으면 True."""
|
||||
try:
|
||||
cols = set(db.conn.get_columns("ws_candles"))
|
||||
for n in ws_candles_all_ema_column_names():
|
||||
if n in cols:
|
||||
return True
|
||||
for n in ("rsi_7", "rsi_14", "rsi_21"):
|
||||
if n in cols:
|
||||
return True
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
Reference in New Issue
Block a user