수정본
This commit is contained in:
323
backend/collect_krx_data.py
Normal file
323
backend/collect_krx_data.py
Normal file
@@ -0,0 +1,323 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user