예제 소스

surge_minute_redis_writer.py

# -*- coding: utf-8 -*-
"""
KIS 모의 분봉만 조회해 Redis 에 적재 (거래/분석 없음).

  · ``surge_paper_trade.py --redis-url ...`` 와 **동일한 watchlist 인자**를 주는 것을 권장합니다.
  · 키: ``hankook:surge:minute:{종목코드}`` JSON (surge_redis_common)

사용
  export KIS_PAPER=1
  export REDIS_URL=redis://127.0.0.1:6379/0
  python surge_minute_redis_writer.py --interval-sec 30 --top-n 150
"""

from __future__ import annotations

import argparse
import os
import sys
import time
from datetime import date, datetime, timedelta

from kis_client import KISClient
from paper_trade import is_trading_day
from surge_monitor import (
    DEFAULTS as SM_DEFAULTS,
    call_with_net_retry,
    fetch_minute_candles,
    heartbeat,
    is_kr_market_hours,
    newline,
    next_market_open,
    resolve_watchlist,
    wait_until_market_open,
)
from surge_redis_common import DEFAULT_MINUTE_KEY_PREFIX, bars_json_from_df, minute_key

try:
    import redis as redis_mod
except ImportError as exc:  # pragma: no cover
    print("redis 패키지가 필요합니다: pip install redis", file=sys.stderr)
    raise SystemExit(1) from exc


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="KIS 1분봉 → Redis 적재 전용")
    p.add_argument(
        "--redis-url",
        type=str,
        default=os.environ.get("REDIS_URL", "redis://127.0.0.1:6379/0"),
        help="미설정·빈 문자열이면 REDIS_URL 또는 redis://127.0.0.1:6379/0",
    )
    p.add_argument("--redis-prefix", type=str, default=DEFAULT_MINUTE_KEY_PREFIX)
    p.add_argument("--redis-ttl-sec", type=int, default=7200, help="키 TTL(초). 장중 덮어쓰기 전제")
    p.add_argument("--codes", type=str, default="")
    p.add_argument("--top-n", type=int, default=SM_DEFAULTS["top_n"])
    p.add_argument("--top-by", type=str, default=SM_DEFAULTS["top_by"], choices=["amount", "marcap", "volume"])
    p.add_argument("--interval-sec", type=int, default=SM_DEFAULTS["interval_sec"])
    p.add_argument("--per-call-sleep", type=float, default=SM_DEFAULTS["per_call_sleep"])
    p.add_argument("--rate-limit-backoff", type=float, default=SM_DEFAULTS["rate_limit_backoff"])
    p.add_argument("--net-retries", type=int, default=SM_DEFAULTS["net_retries"])
    p.add_argument("--net-retry-backoff", type=float, default=SM_DEFAULTS["net_retry_backoff"])
    p.add_argument("--sections", type=str, default=SM_DEFAULTS["sections"])
    p.add_argument("--pages", type=int, default=SM_DEFAULTS["pages"])
    p.add_argument("--min-score", type=int, default=SM_DEFAULTS["min_score"])
    p.add_argument("--max-picks", type=int, default=SM_DEFAULTS["max_picks"])
    p.add_argument("--delay", type=float, default=SM_DEFAULTS["delay"])
    p.add_argument("--force", action="store_true")
    return p.parse_args()


def main() -> int:
    args = parse_args()
    # ``--redis-url "$REDIS_URL"`` 처럼 빈 환경변수를 넘기면 argparse 기본값이 무시되므로 보정
    if not (args.redis_url or "").strip():
        args.redis_url = os.environ.get("REDIS_URL", "redis://127.0.0.1:6379/0")
    if not (args.redis_url or "").strip():
        print("--redis-url 또는 REDIS_URL 을 설정하세요.", file=sys.stderr)
        return 1

    r = redis_mod.Redis.from_url(args.redis_url, decode_responses=True)
    try:
        r.ping()
    except Exception as exc:
        print(f"Redis 연결 실패: {exc}", file=sys.stderr)
        return 1

    watchlist = resolve_watchlist(args)
    if not watchlist:
        print("watchlist 가 비었습니다.", file=sys.stderr)
        return 1

    client = KISClient(paper=True)
    print("[init] KIS 모의 토큰 발급...")
    client.issue_access_token()
    print(f"[init] Redis {args.redis_url}  prefix={args.redis_prefix!r}  종목 {len(watchlist)}")
    print()

    cycle = 0
    try:
        while True:
            now0 = datetime.now()
            if not (is_trading_day(now0.date()) and is_kr_market_hours(now0)):
                wait_until_market_open(len(watchlist))
                newline()
                print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] 장 개시 — Redis 적재 재개")
                continue

            cycle += 1
            t0 = datetime.now()
            ok = 0
            for code, name in watchlist:
                try:
                    df = call_with_net_retry(
                        fetch_minute_candles,
                        client,
                        code,
                        retries=args.net_retries,
                        backoff=args.net_retry_backoff,
                        label=f"{code} {name}",
                    )
                except Exception as exc:
                    newline()
                    print(f"  · {code} 실패: {exc}", file=sys.stderr)
                    if args.per_call_sleep:
                        time.sleep(args.per_call_sleep)
                    continue
                if args.per_call_sleep:
                    time.sleep(args.per_call_sleep)
                if df is None or df.empty:
                    continue
                key = minute_key(code, prefix=args.redis_prefix)
                r.set(key, bars_json_from_df(df), ex=args.redis_ttl_sec)
                ok += 1

            dur = (datetime.now() - t0).total_seconds()
            remain = max(0.0, args.interval_sec - dur)
            target = datetime.now() + timedelta(seconds=remain)
            while True:
                left = (target - datetime.now()).total_seconds()
                if left <= 0:
                    break
                heartbeat(
                    f"[{datetime.now():%H:%M:%S}] Redis writer  cycle#{cycle}  "
                    f"ok={ok}/{len(watchlist)}  next {int(left)}s"
                )
                time.sleep(min(1.0, left))
    except KeyboardInterrupt:
        newline()
        print("[stop] writer 종료")
        return 0


if __name__ == "__main__":
    sys.exit(main())