""" KRX 히스토리컬 데이터 수집 스크립트 pykrx를 사용하여 KOSPI/KOSDAQ 종목의 일봉 OHLCV 데이터를 DB에 저장합니다. 사용법: python collect_krx_data.py # 전체 수집 (10년, 주요 200종목) python collect_krx_data.py --limit 5 --start-date 20240101 # 테스트 (5종목, 2024년~) python collect_krx_data.py --update # 증분 업데이트 (마지막 저장일 이후) python collect_krx_data.py --market kospi # KOSPI만 """ import argparse import asyncio import logging import sys import time from datetime import date, datetime, timedelta from pathlib import Path import numpy as np import pandas as pd from pykrx import stock as pykrx_stock # 프로젝트 루트를 Python 경로에 추가 sys.path.insert(0, str(Path(__file__).parent)) from sqlalchemy import text from sqlalchemy.dialects.postgresql import insert as pg_insert from app.core.database import AsyncSessionLocal, engine, Base from app.models.models import StockDailyPrice logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) logger = logging.getLogger(__name__) # ETF/특수 종목 제외 키워드 SKIP_KEYWORDS = [ "ETF", "레버리지", "인버스", "KODEX", "TIGER", "ARIRANG", "KBSTAR", "HANARO", "SOL", "ACE", "KOSEF", "KINDEX", "선물", "스팩", "리츠", "SPAC", ] def get_stock_list(market: str = "all", limit: int = 0) -> list: """pykrx로 종목 목록 조회 Args: market: "kospi", "kosdaq", "all" limit: 0이면 전체, N이면 상위 N개 Returns: [(stock_code, stock_name), ...] """ today = datetime.now().strftime("%Y%m%d") stocks = [] markets = [] if market in ("kospi", "all"): markets.append("KOSPI") if market in ("kosdaq", "all"): markets.append("KOSDAQ") for mkt in markets: try: tickers = pykrx_stock.get_market_ticker_list(today, market=mkt) for ticker in tickers: name = pykrx_stock.get_market_ticker_name(ticker) if any(k in name for k in SKIP_KEYWORDS): continue stocks.append((ticker, name)) logger.info(f"{mkt}: {len(tickers)}개 조회, 필터링 후 {len(stocks)}개") except Exception as e: logger.error(f"{mkt} 종목 목록 조회 실패: {e}") # 중복 제거 seen = set() unique = [] for code, name in stocks: if code not in seen: seen.add(code) unique.append((code, name)) if limit > 0: unique = unique[:limit] logger.info(f"총 {len(unique)}개 종목 선정") return unique def fetch_ohlcv(stock_code: str, start_date: str, end_date: str) -> pd.DataFrame: """pykrx로 일봉 OHLCV 데이터 조회 Returns: DataFrame with columns: date, open, high, low, close, volume, change_rate, market_cap """ try: df = pykrx_stock.get_market_ohlcv_by_date( start_date, end_date, stock_code, adjusted=True ) if df.empty: return pd.DataFrame() # 시가총액 조회 try: cap_df = pykrx_stock.get_market_cap_by_date(start_date, end_date, stock_code) if not cap_df.empty: df["시가총액"] = cap_df["시가총액"] except Exception: df["시가총액"] = None df = df.reset_index() df.columns = [c.replace("날짜", "date") for c in df.columns] result = pd.DataFrame({ "date": df.iloc[:, 0], # 날짜 "open": df["시가"].astype(int), "high": df["고가"].astype(int), "low": df["저가"].astype(int), "close": df["종가"].astype(int), "volume": df["거래량"].astype(int), "change_rate": df["등락률"] / 100.0 if "등락률" in df.columns else 0, "market_cap": df.get("시가총액"), }) # 거래량 0인 날 제거 (거래정지일) result = result[result["volume"] > 0].reset_index(drop=True) return result except Exception as e: logger.error(f"{stock_code} OHLCV 조회 실패: {e}") return pd.DataFrame() async def create_tables(): """DB 테이블 생성""" async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) logger.info("DB 테이블 확인/생성 완료") async def get_last_trade_date(session, stock_code: str) -> date | None: """특정 종목의 DB 내 마지막 거래일 조회""" result = await session.execute( text( "SELECT MAX(trade_date) FROM stock_daily_prices WHERE stock_code = :code" ), {"code": stock_code}, ) row = result.scalar() return row async def upsert_daily_prices(session, stock_code: str, stock_name: str, df: pd.DataFrame): """DB에 일봉 데이터 upsert (중복 시 업데이트)""" if df.empty: return 0 records = [] for _, row in df.iterrows(): trade_date = row["date"] if isinstance(trade_date, pd.Timestamp): trade_date = trade_date.date() record = { "stock_code": stock_code, "stock_name": stock_name, "trade_date": trade_date, "open_price": int(row["open"]), "high_price": int(row["high"]), "low_price": int(row["low"]), "close_price": int(row["close"]), "volume": int(row["volume"]), "change_rate": float(row["change_rate"]) if pd.notna(row["change_rate"]) else None, "market_cap": int(row["market_cap"]) if pd.notna(row.get("market_cap")) else None, } records.append(record) # PostgreSQL ON CONFLICT upsert (batch) batch_size = 500 inserted = 0 for i in range(0, len(records), batch_size): batch = records[i:i + batch_size] stmt = pg_insert(StockDailyPrice).values(batch) stmt = stmt.on_conflict_do_update( constraint="uq_stock_daily", set_={ "stock_name": stmt.excluded.stock_name, "open_price": stmt.excluded.open_price, "high_price": stmt.excluded.high_price, "low_price": stmt.excluded.low_price, "close_price": stmt.excluded.close_price, "volume": stmt.excluded.volume, "change_rate": stmt.excluded.change_rate, "market_cap": stmt.excluded.market_cap, }, ) await session.execute(stmt) inserted += len(batch) await session.commit() return inserted async def collect_all( market: str = "all", limit: int = 0, start_date: str = "20160101", update_mode: bool = False, ): """메인 수집 로직""" await create_tables() end_date = datetime.now().strftime("%Y%m%d") # 종목 목록 stock_list = get_stock_list(market=market, limit=limit) if not stock_list: logger.error("종목 목록이 비어있습니다") return total = len(stock_list) total_inserted = 0 success_count = 0 fail_count = 0 async with AsyncSessionLocal() as session: for i, (stock_code, stock_name) in enumerate(stock_list): try: actual_start = start_date if update_mode: last_date = await get_last_trade_date(session, stock_code) if last_date: # 마지막 저장일 다음 날부터 next_day = last_date + timedelta(days=1) actual_start = next_day.strftime("%Y%m%d") if next_day >= date.today(): logger.info(f"[{i+1}/{total}] {stock_code} ({stock_name}): 이미 최신") success_count += 1 continue logger.info( f"[{i+1}/{total}] {stock_code} ({stock_name}): " f"{actual_start} ~ {end_date} 수집 중..." ) df = fetch_ohlcv(stock_code, actual_start, end_date) if df.empty: logger.warning(f" {stock_code}: 데이터 없음") fail_count += 1 continue count = await upsert_daily_prices(session, stock_code, stock_name, df) total_inserted += count success_count += 1 logger.info(f" {stock_code}: {count}개 행 저장 (총 {len(df)}일)") except Exception as e: logger.error(f" {stock_code} ({stock_name}) 실패: {e}") fail_count += 1 # 에러 시에도 계속 진행 try: await session.rollback() except Exception: pass # pykrx rate limit 방지 time.sleep(0.5) logger.info("=" * 60) logger.info(f"수집 완료!") logger.info(f" 성공: {success_count}/{total}") logger.info(f" 실패: {fail_count}/{total}") logger.info(f" 총 저장 행: {total_inserted:,}") logger.info("=" * 60) def main(): parser = argparse.ArgumentParser(description="KRX 히스토리컬 데이터 수집") parser.add_argument( "--market", choices=["kospi", "kosdaq", "all"], default="all", help="수집 시장 (기본: all)", ) parser.add_argument( "--limit", type=int, default=0, help="종목 수 제한 (0=전체, 테스트용)" ) parser.add_argument( "--start-date", default="20160101", help="수집 시작일 YYYYMMDD (기본: 20160101)", ) parser.add_argument( "--update", action="store_true", help="증분 업데이트 모드 (마지막 저장일 이후만 수집)", ) args = parser.parse_args() logger.info(f"KRX 데이터 수집 시작") logger.info(f" 시장: {args.market}") logger.info(f" 종목 제한: {'전체' if args.limit == 0 else args.limit}") logger.info(f" 시작일: {args.start_date}") logger.info(f" 증분 모드: {args.update}") start_time = time.time() asyncio.run(collect_all( market=args.market, limit=args.limit, start_date=args.start_date, update_mode=args.update, )) elapsed = time.time() - start_time logger.info(f"총 소요시간: {elapsed:.1f}초 ({elapsed/60:.1f}분)") if __name__ == "__main__": main()