324 lines
10 KiB
Python
324 lines
10 KiB
Python
"""
|
|
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()
|